Skip to content

Wrong import for batch_generate in MLXLM #1770

@samwestb

Description

@samwestb

Describe the issue as clearly as possible:

Line 176 in mlxlm.py has from mlx_lm import generate_batch when it should be batch_generate

However, after fixing the import issue and running the same code snippet again I found that the BatchGenerator in mlx-lm doesn't accept the logits_processors kwarg.

TypeError: BatchGenerator.__init__() got an unexpected keyword argument 'logits_processors'

Steps/code to reproduce the bug:

import outlines
import mlx_lm

# Load the model
model = outlines.from_mlxlm(
    *mlx_lm.load("mlx-community/granite-4.0-h-micro-4bit")
)

# Generate text in batches
result = model.batch(["What's the capital of Lithuania?", "What's the capital of Latvia?"], max_tokens=20)
print(result) # ['Vilnius', 'Riga']

Expected result:

The result of the batch generation.

Error message:

Fetching 10 files: 100%|█████████████████████████████████████████████████| 10/10 [00:00<00:00, 16156.80it/s]
Traceback (most recent call last):
  File "/Users/sam/testing/test.py", line 10, in <module>
    result = model.batch(["What's the capital of Lithuania?", "What's the capital of Latvia?"], max_tokens=2
0)                                                                                                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
^^                                                                                                            File "/Users/sam/testing/.venv/lib/python3.12/site-packages/outlines/models/base.py", line 168, in batch
    return generator.batch(model_input, **inference_kwargs) # type: ignore
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sam/testing/.venv/lib/python3.12/site-packages/outlines/generator.py", line 319, in batch
    return self.model.generate_batch(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/sam/testing/.venv/lib/python3.12/site-packages/outlines/models/mlxlm.py", line 176, in genera
te_batch                                                                                                        from mlx_lm import generate_batch
ImportError: cannot import name 'generate_batch' from 'mlx_lm' (/Users/sam/testing/.venv/lib/python3.12/site
-packages/mlx_lm/__init__.py)

Outlines/Python version information:

Version information

``` 1.2.7 Python 3.12.11 (main, Sep 18 2025, 19:41:45) [Clang 20.1.4 ] annotated-types==0.7.0 attrs==25.4.0 certifi==2025.10.5 charset-normalizer==3.4.4 cloudpickle==3.1.1 diskcache==5.6.3 filelock==3.20.0 fsspec==2025.9.0 genson==1.3.0 hf-xet==1.1.10 huggingface-hub==0.35.3 idna==3.11 jinja2==3.1.6 jsonpath-ng==1.7.0 jsonschema==4.25.1 jsonschema-specifications==2025.9.1 markupsafe==3.0.3 mlx==0.29.2 mlx-lm==0.28.2 mlx-metal==0.29.2 numpy==2.3.3 outlines==1.2.7 outlines-core==0.2.11 packaging==25.0 pillow==11.3.0 ply==3.11 protobuf==6.32.1 pydantic==2.12.2 pydantic-core==2.41.4 pyyaml==6.0.3 referencing==0.37.0 regex==2025.9.18 requests==2.32.5 rpds-py==0.27.1 safetensors==0.6.2 tokenizers==0.22.1 tqdm==4.67.1 transformers==4.57.1 typing-extensions==4.15.0 typing-inspection==0.4.2 urllib3==2.5.0
</details>


### Context for the issue:

_No response_

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions