Clamp out of range shift amounts

[ayaangazali]

On the CPU, left_shift returns a different answer for the same values depending on how long the array is:

>>> for n in (1, 2, 3, 4, 7, 8, 1000):
...     x = mx.ones(n, mx.int32)
...     print(n, mx.left_shift(x, mx.full((n,), 32, mx.int32)).tolist()[:8])
1    [1]
2    [1, 1]
3    [1, 1, 1]
4    [0, 0, 0, 0]
7    [0, 0, 0, 0, 1, 1, 1]
8    [0, 0, 0, 0, 0, 0, 0, 0]
1000 [0, 0, 0, ...]

At n=7 a single array contains both answers: the vectorized body produces 0 and the scalar tail produces 1. A strided view takes the scalar path too, so mx.left_shift(x[::2], 32) disagrees with mx.left_shift(x, 32).

The cause is that a shift amount that is negative or at least the operand width is undefined in C++, and the two paths landed on different answers for it. binary_ops.h maps both ops straight onto the built in operators:

DEFAULT_BINARY_OP(LeftShift, operator<<)
DEFAULT_BINARY_OP(RightShift, operator>>)

The scalar overload gets the hardware behaviour, which on arm64 and x86 masks the count against the width, so 1 << 32 becomes 1 << 0. The vectorized overload goes through the Accelerate vector types, which saturate to 0 instead.

numpy and pytorch agree with each other on every case, and mlx's vectorized path already matches them:

   x     sh      numpy <<   torch <<      numpy >>   torch >>
   1     32             0          0             0          0
  16     32             0          0             0          0
 -16     32             0          0            -1         -1
   1     -1             0          0             0          0
 -16     -1             0          0            -1         -1

So the rule is: out of range means zero for a left shift, and the sign bit for a right shift. This clamps the amount to make that true on both paths. The mask keeps the shift itself in range so the operation is no longer undefined, and the select picks the out of range answer.

Verified against numpy over all 8 integer dtypes, array lengths 1, 3, 5, 7, 8, 16, 33, 64 and 257, every shift amount from 0 to width+2 plus 2*width and 8*width and the negative amounts, and the type min, max, and assorted values. That is 39744 combinations, all matching, where main misses 2697 of them. Also checked 9999 element arrays of randomly mixed amounts, which exercises the vector body and the scalar tail in a single call.

The extra select is not free but it is small, 8M elements on an M4:

            main     this
int32  <<   1.045    1.206 ms
int32  >>   1.007    1.271 ms
uint8  <<   0.254    0.274 ms
uint8  >>   0.228    0.292 ms
int64  <<   2.126    2.240 ms
int64  >>   2.388    2.249 ms

This started as a CPU only change, because Metal and CUDA also use a bare x << y and I have no GPU here to see what they actually returned. CI answered that: the macOS run failed on the new test with left_shift mlx.core.uint32 n=1 1 by 32, which is the GPU masking the count against the width exactly like the CPU scalar path did. So the same clamp is now in metal/kernels/binary_ops.h and cuda/device/binary_ops.cuh, and all three backends agree with numpy and pytorch.

The GPU side is the same shape as the CPU one, with the signed check behind if constexpr so the unsigned instantiations do not emit a tautological comparison under -DCMAKE_COMPILE_WARNING_AS_ERROR=ON. Shifts are only instantiated for the eight integer types, so there is no bool, float or complex path to worry about. I still cannot build Metal or CUDA locally, so CI is the check on those two.

test_ops.py, test_array.py, test_reduce.py, test_nn.py, test_linalg.py, test_fft.py, test_compile.py, test_vmap.py, test_autograd.py, test_random.py, test_double.py and the C++ suite (249 cases, 3350 assertions) pass. The added test fails on main with left_shift mlx.core.uint32 n=1 1 by 32.

I am a freshman working through this codebase and I used Claude Code alongside it. Every number above came from a run on this machine.

[ayaangazali]

CI caught a real gap in this and I have pushed the fix.

The macOS job failed on the test I added:

FAIL: test_bitwise_ops (test_ops.TestOps.test_bitwise_ops)
AssertionError: False is not true : left_shift mlx.core.uint32 n=1 1 by 32

That is the GPU, not the CPU. It is the same masking behaviour the CPU scalar path had, so a CPU only change left mx.left_shift(x, 32) returning 1 on Metal and 0 on the CPU. That is a worse state than before, since the answer would then depend on the device rather than on the array length.

So the clamp now also lives in metal/kernels/binary_ops.h and cuda/device/binary_ops.cuh. The signed check sits behind if constexpr so the unsigned instantiations do not trip -DCMAKE_COMPILE_WARNING_AS_ERROR=ON, and shifts are only instantiated for the eight integer types so there is no bool, float or complex path involved.

Worth saying plainly: I could not have found this locally, my build here is CPU only. The failing job is what told me what the GPU actually does, and it is the check on the two kernels I cannot compile.

[ayaangazali]

The GPU fix works. The test that was failing now passes on the Metal run:

test_bitwise_ops (test_ops.TestOps.test_bitwise_ops) ... ok

Ran 840 tests in 114.049s

OK (skipped=50)

The three macOS (*, metal) builds and the CUDA builds are all green too, so the kernel changes compile on both backends.

Test macOS is still red, but on a different and unrelated step. After the suite passes, the job dies downloading a wheel:

##[group]Build example extension
error: Request failed after 3 retries in 7.0s
  Caused by: Failed to fetch: `https://files.pythonhosted.org/packages/.../nanobind-2.13.0-py3-none-any.whl.metadata`
  Caused by: HTTP status server error (502 Bad Gateway) for url (...)

That is the same PyPI 502 that is failing the job on my other PR, and my four other open PRs have this job green, so it is a transient registry problem rather than anything in this branch. A re-run should clear it. I cannot trigger one here.