Clamp out of range shift amounts

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 <int N, typename T>
+Simd<bool, N> shift_in_range(Simd<T, N> y) {
+  auto in_range = y < Simd<T, N>(sizeof(T) * 8);
+  if constexpr (std::is_signed_v<T>) {
+    in_range = in_range && y >= Simd<T, N>(0);
+  }
+  return in_range;
+}
+
+struct LeftShift {
+  template <int N, typename T>
+  Simd<T, N> operator()(Simd<T, N> x, Simd<T, N> y) {
+    Simd<T, N> amount = y & Simd<T, N>(sizeof(T) * 8 - 1);
+    Simd<T, N> shifted = x << amount;
+    return select(shift_in_range(y), shifted, Simd<T, N>(0));
+  }
+  BINARY_SINGLE()
+};
+
+struct RightShift {
+  template <int N, typename T>
+  Simd<T, N> operator()(Simd<T, N> x, Simd<T, N> y) {
+    auto in_range = shift_in_range(y);
+    Simd<T, N> amount = select(in_range, y, Simd<T, N>(sizeof(T) * 8 - 1));
+    Simd<T, N> shifted = x >> amount;
+    if constexpr (std::is_signed_v<T>) {
+      // Shifting a negative value past the width leaves the sign bit behind.
+      return shifted;
+    } else {
+      return select(in_range, shifted, Simd<T, N>(0));
+    }
+  }
+  BINARY_SINGLE()
+};
+
 DEFAULT_BINARY_OP(Remainder, remainder)
 DEFAULT_BINARY_OP(Maximum, maximum)
 DEFAULT_BINARY_OP(Minimum, minimum)

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 <typename T>
+__device__ bool shift_in_range(T y) {
+  bool in_range = y < static_cast<T>(sizeof(T) * 8);
+  if constexpr (cuda::std::is_signed_v<T>) {
+    in_range = in_range && y >= 0;
+  }
+  return in_range;
+}
+
 struct LeftShift {
   template <typename T>
   __device__ T operator()(T x, T y) {
-    return x << y;
+    T amount = y & static_cast<T>(sizeof(T) * 8 - 1);
+    return shift_in_range(y) ? static_cast<T>(x << amount) : static_cast<T>(0);
   };
 };
 
 struct RightShift {
   template <typename T>
   __device__ T operator()(T x, T y) {
-    return x >> y;
+    bool in_range = shift_in_range(y);
+    T amount = in_range ? y : static_cast<T>(sizeof(T) * 8 - 1);
+    T shifted = x >> amount;
+    if constexpr (cuda::std::is_signed_v<T>) {
+      // Shifting a negative value past the width leaves the sign bit behind.
+      return shifted;
+    } else {
+      return in_range ? shifted : static_cast<T>(0);
+    }
   };
 };
 

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 <typename T>
+inline bool shift_in_range(T y) {
+  bool in_range = y < static_cast<T>(sizeof(T) * 8);
+  if constexpr (metal::is_signed_v<T>) {
+    in_range = in_range && y >= 0;
+  }
+  return in_range;
+}
+
 struct LeftShift {
   template <typename T>
   T operator()(T x, T y) thread {
-    return x << y;
+    T amount = y & static_cast<T>(sizeof(T) * 8 - 1);
+    return shift_in_range(y) ? static_cast<T>(x << amount) : static_cast<T>(0);
   };
 };
 
 struct RightShift {
   template <typename T>
   T operator()(T x, T y) thread {
-    return x >> y;
+    bool in_range = shift_in_range(y);
+    T amount = in_range ? y : static_cast<T>(sizeof(T) * 8 - 1);
+    T shifted = x >> amount;
+    if constexpr (metal::is_signed_v<T>) {
+      // Shifting a negative value past the width leaves the sign bit behind.
+      return shifted;
+    } else {
+      return in_range ? shifted : static_cast<T>(0);
+    }
   };
 };
 

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)