From e2e0fe41605a0b3d1e12548c98c180b347a1f21e Mon Sep 17 00:00:00 2001 From: ayaangazali Date: Sat, 15 Aug 2026 10:55:52 -0700 Subject: [PATCH 1/2] Clamp out of range shift amounts on the CPU A shift amount that is negative or at least the operand width is undefined in C++. The scalar path took the hardware answer and the vectorized path saturated, so left_shift on the same values returned different results depending on the length of the array. Clamp the amount so both paths agree and match numpy and pytorch. --- mlx/backend/cpu/binary_ops.h | 42 ++++++++++++++++++++++++++++++++++-- python/tests/test_ops.py | 28 ++++++++++++++++++++++++ 2 files changed, 68 insertions(+), 2 deletions(-) diff --git a/mlx/backend/cpu/binary_ops.h b/mlx/backend/cpu/binary_ops.h index d50751ce3d..c617c70f92 100644 --- a/mlx/backend/cpu/binary_ops.h +++ b/mlx/backend/cpu/binary_ops.h @@ -33,8 +33,46 @@ DEFAULT_BINARY_OP(LogicalOr, operator||) DEFAULT_BINARY_OP(BitwiseAnd, operator&) DEFAULT_BINARY_OP(BitwiseOr, operator|) DEFAULT_BINARY_OP(BitwiseXor, operator^) -DEFAULT_BINARY_OP(LeftShift, operator<<) -DEFAULT_BINARY_OP(RightShift, operator>>) +// A shift amount that is negative or at least the operand width is undefined +// in C++, so the scalar and vectorized paths disagreed about it and the answer +// depended on the array length. Clamp the amount instead, which gives what +// numpy and pytorch return: zero for a left shift, and the sign bit for a +// right shift. +template +Simd shift_in_range(Simd y) { + auto in_range = y < Simd(sizeof(T) * 8); + if constexpr (std::is_signed_v) { + in_range = in_range && y >= Simd(0); + } + return in_range; +} + +struct LeftShift { + template + Simd operator()(Simd x, Simd y) { + Simd amount = y & Simd(sizeof(T) * 8 - 1); + Simd shifted = x << amount; + return select(shift_in_range(y), shifted, Simd(0)); + } + BINARY_SINGLE() +}; + +struct RightShift { + template + Simd operator()(Simd x, Simd y) { + auto in_range = shift_in_range(y); + Simd amount = select(in_range, y, Simd(sizeof(T) * 8 - 1)); + Simd shifted = x >> amount; + if constexpr (std::is_signed_v) { + // Shifting a negative value past the width leaves the sign bit behind. + return shifted; + } else { + return select(in_range, shifted, Simd(0)); + } + } + BINARY_SINGLE() +}; + DEFAULT_BINARY_OP(Remainder, remainder) DEFAULT_BINARY_OP(Maximum, maximum) DEFAULT_BINARY_OP(Minimum, minimum) diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index dd0a93001a..ac308a0ef7 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -3448,6 +3448,34 @@ def test_bitwise_ops(self): out_np = getattr(np, op)(a_np, b_np) self.assertTrue(np.array_equal(np.array(out_mlx), out_np)) + # A shift amount that is negative or at least the operand width is + # undefined in C++, so the scalar and vectorized paths disagreed and + # the answer depended on the length of the array. + for t in types: + nt = np.array(mx.zeros(1, dtype=t)).dtype + info = np.iinfo(nt) + width = t.size * 8 + amounts = [0, 1, width - 1, width, width + 3] + values = [1, 16, info.max] + if info.min < 0: + amounts += [-1, -width] + values += [-16, info.min] + for n in (1, 7, 8, 100): + for v in values: + for sh in amounts: + a_np = np.full(n, v, nt) + b_np = np.full(n, sh, nt) + a_mlx = mx.array(a_np) + b_mlx = mx.array(b_np) + for op in ("left_shift", "right_shift"): + self.assertTrue( + np.array_equal( + np.array(getattr(mx, op)(a_mlx, b_mlx)), + getattr(np, op)(a_np, b_np), + ), + msg=f"{op} {t} n={n} {v} by {sh}", + ) + for t in types: a_mlx = a.astype(t) a_np = np.array(a_mlx) From 2f99a1eb1361f2c7b8a27078bea59cae9e7259fa Mon Sep 17 00:00:00 2001 From: ayaangazali Date: Sat, 15 Aug 2026 23:15:08 -0700 Subject: [PATCH 2/2] Clamp out of range shift amounts on Metal and CUDA too CI showed the GPU masks the shift count against the operand width, so a CPU only fix left left_shift(1, 32) returning 1 there and 0 on the CPU. Apply the same clamp in both GPU kernels. --- mlx/backend/cuda/device/binary_ops.cuh | 25 +++++++++++++++++++++++-- mlx/backend/metal/kernels/binary_ops.h | 25 +++++++++++++++++++++++-- 2 files changed, 46 insertions(+), 4 deletions(-) diff --git a/mlx/backend/cuda/device/binary_ops.cuh b/mlx/backend/cuda/device/binary_ops.cuh index b0b7962807..75111b0eba 100644 --- a/mlx/backend/cuda/device/binary_ops.cuh +++ b/mlx/backend/cuda/device/binary_ops.cuh @@ -269,17 +269,38 @@ struct BitwiseXor { }; }; +// A shift amount that is negative or at least the operand width is undefined, +// so clamp it to what numpy and pytorch return: zero for a left shift, and the +// sign bit for a right shift. +template +__device__ bool shift_in_range(T y) { + bool in_range = y < static_cast(sizeof(T) * 8); + if constexpr (cuda::std::is_signed_v) { + in_range = in_range && y >= 0; + } + return in_range; +} + struct LeftShift { template __device__ T operator()(T x, T y) { - return x << y; + T amount = y & static_cast(sizeof(T) * 8 - 1); + return shift_in_range(y) ? static_cast(x << amount) : static_cast(0); }; }; struct RightShift { template __device__ T operator()(T x, T y) { - return x >> y; + bool in_range = shift_in_range(y); + T amount = in_range ? y : static_cast(sizeof(T) * 8 - 1); + T shifted = x >> amount; + if constexpr (cuda::std::is_signed_v) { + // Shifting a negative value past the width leaves the sign bit behind. + return shifted; + } else { + return in_range ? shifted : static_cast(0); + } }; }; diff --git a/mlx/backend/metal/kernels/binary_ops.h b/mlx/backend/metal/kernels/binary_ops.h index 863d6369e2..0bcee104ce 100644 --- a/mlx/backend/metal/kernels/binary_ops.h +++ b/mlx/backend/metal/kernels/binary_ops.h @@ -303,17 +303,38 @@ struct BitwiseXor { }; }; +// A shift amount that is negative or at least the operand width is undefined, +// so clamp it to what numpy and pytorch return: zero for a left shift, and the +// sign bit for a right shift. +template +inline bool shift_in_range(T y) { + bool in_range = y < static_cast(sizeof(T) * 8); + if constexpr (metal::is_signed_v) { + in_range = in_range && y >= 0; + } + return in_range; +} + struct LeftShift { template T operator()(T x, T y) thread { - return x << y; + T amount = y & static_cast(sizeof(T) * 8 - 1); + return shift_in_range(y) ? static_cast(x << amount) : static_cast(0); }; }; struct RightShift { template T operator()(T x, T y) thread { - return x >> y; + bool in_range = shift_in_range(y); + T amount = in_range ? y : static_cast(sizeof(T) * 8 - 1); + T shifted = x >> amount; + if constexpr (metal::is_signed_v) { + // Shifting a negative value past the width leaves the sign bit behind. + return shifted; + } else { + return in_range ? shifted : static_cast(0); + } }; };