Skip to content

fix: accept PyTorch tokenizer inputs in MLX generate - #363

Open
Vaishnavi220506 wants to merge 1 commit into
lyogavin:mainfrom
Vaishnavi220506:codex/fix-mlx-input-types
Open

Vaishnavi220506 wants to merge 1 commit into
lyogavin:mainfrom
Vaishnavi220506:codex/fix-mlx-input-types

Conversation

@Vaishnavi220506

Copy link
Copy Markdown

Problem

On macOS, AutoModel.from_pretrained() routes models through AirLLMLlamaMlx. generate() passed the PyTorch tensor returned by the standard tokenizer directly to mlx.nn.Embedding, which rejects it before generation starts.

Changes

  • Normalize input IDs at the generate() boundary.
  • Preserve existing MLX arrays without copying them.
  • Convert PyTorch tensors through a detached CPU NumPy view before creating the MLX array.
  • Add regression coverage for PyTorch tensors, MLX arrays, and Python sequences.

Validation

  • python -m compileall -q air_llm/airllm/airllm_llama_mlx.py air_llm/tests/test_mlx_input_types.py
  • Targeted repository tests: 47 passed, 3 MLX tests skipped on Windows.
  • python -m build
  • python -m twine check dist/*

Closes #280

@Vaishnavi220506
Vaishnavi220506 force-pushed the codex/fix-mlx-input-types branch from 816348b to be9c45e Compare September 12, 2026 11:17

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

AirLLMLlamaMlx.generate() crashes on torch.LongTensor input — "Cannot index mlx array using the given type"

1 participant