fix(processors): NumpyTensorAdapter.apply_mask doesn't broadcast a lower-rank mask like torch's masked_fill

[chuenchen309]

Problem

apply_mask used boolean fancy-index assignment (result[mask] = value), which requires the mask's shape to exactly match the tensor's shape — numpy raises IndexError instead of broadcasting when the shapes only match under normal broadcasting rules (e.g. a 1D mask applied to a 2D tensor, a common pattern for reusing the same per-token mask across every row of a batch).

TorchTensorAdapter.apply_mask already uses torch.masked_fill, and MLXTensorAdapter already uses mlx.where, both of which broadcast correctly. Only the numpy adapter lacked this, making numpy-backed generation crash on a mask shape that works fine on the other two backends.

Fix

Use numpy.where(mask, value, tensor), matching MLX's approach and torch's masked_fill's broadcasting semantics. Also non-mutating (like masked_fill), unlike the previous copy-then-index-assign — a behavior improvement consistent with the other two adapters.

Testing

• Added test_tensor_adapter_apply_mask_broadcasts_lower_rank_mask (parametrized across all three frameworks) using a 1D mask against a 2D tensor.
• Confirmed red→green: reverting only numpy.py reproduces IndexError: boolean index did not match indexed array along axis 0; reapplying passes.
• Full tests/processors/test_tensor_adapters.py: 24 passed, 12 skipped (MLX unavailable in this environment, pre-existing and unrelated to this change).
• ruff and mypy clean via pre-commit run.
• xref: no open PR touches tensor_adapters/numpy.py.

AI-Generated disclosure

Found via an AI-assisted code review pass (Claude Code) over outlines/src/outlines/processors/tensor_adapters/. I personally compared numpy's implementation against the already-correct torch/mlx siblings, reproduced the broadcasting crash, verified the fix preserves exact-shape-mask behavior while fixing the broadcast case, and ran the relevant test suite plus lint/type checks before submitting.

[ErenAta16]

Confirmed this is a real crash, not just a semantic mismatch: numpy boolean-array indexing (tensor[mask] = value) requires mask's leading dimensions to exactly match tensor's, so a mask of shape (3,) against a tensor of shape (2, 3) raises IndexError rather than broadcasting per-row. np.where(mask, value, tensor) does standard numpy broadcasting instead, which is exactly what torch's masked_fill does and what the per-row banned-token-mask use case needs. New test correctly checks the broadcast case (mask shape (3,) against tensor shape (2, 3)) independent of the existing same-shape test, so there's no regression risk on the already-covered path. LGTM.

[RobinPicard]

Good improvement, thanks!

[github-actions]

📚 Documentation preview: https://dottxt-ai.github.io/outlines/pr-preview/pr-1927/

Preview updates automatically with each commit.