fix: divmod should floor the quotient to match numpy (#4119)

mlx/backend/cpu/binary.cpp

@@ -2,6 +2,7 @@
 
 #include <cassert>
 #include <cmath>
+#include <limits>
 #include <sstream>
 #include <type_traits>
 
@@ -46,11 +47,41 @@ void DivMod::eval_cpu(
                     out_a = array::unsafe_weak_copy(out_a),
                     out_b = array::unsafe_weak_copy(out_b),
                     bopt]() mutable {
+    // Match numpy's semantics: the quotient is floor(a / b) and the remainder
+    // carries the divisor's sign, so q * b + r == a holds for every sign
+    // combination.
     auto integral_op = [](auto x, auto y) {
-      return std::make_pair(x / y, x % y);
+      // Start from the truncating quotient/remainder and shift both down by
+      // one when the signs differ. This avoids (a - r), which can overflow
+      // signed integers at the extremes (e.g. INT_MAX / -2).
+      auto q = x / y;
+      auto r = x % y;
+      if (r != 0 && ((r < 0) != (y < 0))) {
+        q -= 1;
+        r += y;
+      }
+      return std::make_pair(q, r);
     };
     auto float_op = [](auto x, auto y) {
-      return std::make_pair(std::trunc(x / y), std::fmod(x, y));
+      auto r = std::fmod(x, y);
+      decltype(r) q;
+      if (y == 0) {
+        // numpy treats b == 0 specially: the quotient is a / b (which yields
+        // +/-inf for a nonzero) and the remainder is fmod(a, b) (nan).
+        q = static_cast<decltype(r)>(x / y);
+      } else if (std::isnan(x) || std::isnan(y) || std::isinf(x)) {
+        // numpy's floor_divide returns nan for nan inputs and for an infinite
+        // dividend (but keeps +/-inf when a finite dividend overflows).
+        q = std::numeric_limits<decltype(r)>::quiet_NaN();
+      } else {
+        // floor(a / b) matches numpy bit for bit; deriving the quotient from
+        // the remainder instead can differ by one ulp.
+        q = std::floor(x / y);
+      }
+      if (r != 0 && ((r < 0) != (y < 0))) {
+        r += y;
+      }
+      return std::make_pair(q, r);
     };
 
     dispatch_all_types(out_a.dtype(), [&](auto type_tag) {

mlx/backend/cuda/device/binary_ops.cuh

@@ -13,17 +13,6 @@ struct Add {
   }
 };
 
-struct FloorDivide {
-  template <typename T>
-  __device__ T operator()(T x, T y) {
-    if constexpr (cuda::std::is_integral_v<T>) {
-      return x / y;
-    } else {
-      return cuda::std::trunc(x / y);
-    }
-  }
-};
-
 struct Divide {
   template <typename T>
   __device__ T operator()(T x, T y) {
@@ -293,7 +282,43 @@ struct ArcTan2 {
 struct DivMod {
   template <typename T>
   __device__ cuda::std::array<T, 2> operator()(T x, T y) {
-    return {FloorDivide{}(x, y), Remainder{}(x, y)};
+    if constexpr (cuda::std::is_integral_v<T>) {
+      // Integer floor-divmod without overflow: start from the truncating
+      // quotient/remainder and shift both down by one when the signs differ.
+      // This avoids computing (a - r), which can overflow signed integers at
+      // the extremes (e.g. INT_MAX / -2).
+      auto q = x / y;
+      auto r = x % y;
+      if (r != 0 && ((r < 0) != (y < 0))) {
+        q -= 1;
+        r += y;
+      }
+      return {q, r};
+    } else if constexpr (is_complex_v<T>) {
+      auto r = Remainder{}(x, y);
+      return {(x - r) / y, r};
+    } else {
+      // numpy semantics: the quotient is floor(a / b) and the remainder
+      // carries the divisor's sign, so q * b + r == a holds for every sign
+      // combination. b == 0 yields a / b; nan inputs and an infinite dividend
+      // yield nan, matching numpy's floor_divide. floor(a / b) matches numpy
+      // bit for bit (deriving the quotient from the remainder can differ by
+      // one ulp).
+      auto r = cuda::std::fmod(x, y);
+      T q;
+      if (y == 0) {
+        q = x / y;
+      } else if (cuda::std::isnan(x) || cuda::std::isnan(y) ||
+                 cuda::std::isinf(x)) {
+        q = cuda::std::numeric_limits<T>::quiet_NaN();
+      } else {
+        q = cuda::std::floor(x / y);
+      }
+      if (r != 0 && ((r < 0) != (y < 0))) {
+        r += y;
+      }
+      return {q, r};
+    }
   };
 };
 

mlx/backend/metal/kernels/binary_ops.h

@@ -14,25 +14,6 @@ struct Add {
   }
 };
 
-struct FloorDivide {
-  template <typename T>
-  T operator()(T x, T y) thread {
-    return x / y;
-  }
-  template <>
-  float operator()(float x, float y) thread {
-    return trunc(x / y);
-  }
-  template <>
-  half operator()(half x, half y) thread {
-    return trunc(x / y);
-  }
-  template <>
-  bfloat16_t operator()(bfloat16_t x, bfloat16_t y) thread {
-    return trunc(x / y);
-  }
-};
-
 struct Divide {
   template <typename T>
   T operator()(T x, T y) thread {
@@ -326,7 +307,47 @@ struct ArcTan2 {
 
 struct DivMod {
   template <typename T>
-  metal::array<T, 2> operator()(T x, T y) thread {
-    return {FloorDivide{}(x, y), Remainder{}(x, y)};
+  metal::enable_if_t<metal::is_integral_v<T>, metal::array<T, 2>> operator()(
+      T x,
+      T y) thread {
+    // Integer floor-divmod without overflow: start from the truncating
+    // quotient/remainder and shift both down by one when the signs differ.
+    // This avoids computing (a - r), which can overflow signed integers at
+    // the extremes (e.g. INT_MAX / -2).
+    auto q = x / y;
+    auto r = x % y;
+    if (r != 0 && ((r < 0) != (y < 0))) {
+      q -= 1;
+      r += y;
+    }
+    return {q, r};
+  }
+  template <typename T>
+  metal::enable_if_t<!metal::is_integral_v<T>, metal::array<T, 2>> operator()(
+      T x,
+      T y) thread {
+    // numpy semantics: the quotient is floor(a / b) and the remainder carries
+    // the divisor's sign, so q * b + r == a holds for every sign combination.
+    // b == 0 yields a / b; nan inputs and an infinite dividend yield nan,
+    // matching numpy's floor_divide. floor(a / b) matches numpy bit for bit
+    // (deriving the quotient from the remainder can differ by one ulp).
+    T r = fmod(x, y);
+    T q;
+    if (y == 0) {
+      q = x / y;
+    } else if (x != x || y != y || (T(1) / x == T(0))) {
+      q = T(0) / T(0);  // nan
+    } else {
+      q = floor(x / y);
+    }
+    if (r != 0 && ((r < 0) != (y < 0))) {
+      r += y;
+    }
+    return {q, r};
+  }
+  template <>
+  metal::array<complex64_t, 2> operator()(complex64_t x, complex64_t y) thread {
+    auto r = Remainder{}(x, y);
+    return {(x - r) / y, r};
   };
 };

python/tests/test_ops.py

@@ -3217,13 +3217,45 @@ def test_divmod(self):
         types = [np.uint16, np.uint32, np.int32, np.float16, np.float32]
         for s1, s2 in sizes:
             for t in types:
-                a_np = np.random.uniform(1, 100, size=s1).astype(t)
-                b_np = np.random.uniform(1, 100, size=s2).astype(t)
+                if np.issubdtype(t, np.unsignedinteger):
+                    a_np = np.random.uniform(1, 100, size=s1).astype(t)
+                    b_np = np.random.uniform(1, 100, size=s2).astype(t)
+                else:
+                    # Cover negative operands: the quotient must floor toward
+                    # -inf and the remainder must carry the divisor's sign.
+                    a_np = np.random.uniform(-100, 100, size=s1).astype(t)
+                    b_np = np.random.uniform(-100, 100, size=s2).astype(t)
+                    b_np[b_np == 0] = 1
                 np_out = np.divmod(a_np, b_np)
                 mx_out = mx.divmod(mx.array(a_np), mx.array(b_np))
                 self.assertTrue(
                     np.allclose(np_out[0], mx_out[0]), msg=f"Shapes {s1} {s2}, Type {t}"
                 )
+                self.assertTrue(
+                    np.allclose(np_out[1], mx_out[1]), msg=f"Shapes {s1} {s2}, Type {t}"
+                )
+
+        # Regression: negative operands used to truncate the quotient toward
+        # zero, which broke q * b + r == a and disagreed with numpy.
+        for a_np, b_np in [(-7, 2), (7, -2), (-7, -2), (7, 2)]:
+            for t in [np.int32, np.float32]:
+                a = np.array(a_np, dtype=t)
+                b = np.array(b_np, dtype=t)
+                q, r = mx.divmod(mx.array(a), mx.array(b))
+                np_q, np_r = np.divmod(a, b)
+                self.assertEqual(np.array(q).item(), np_q.item())
+                self.assertEqual(np.array(r).item(), np_r.item())
+
+        # Signed integer extremes: the naive (a - r) construction overflows
+        # here (e.g. INT_MAX - (-1)), so the quotient must come from shifting
+        # the truncating result instead.
+        for t, maxv in [(np.int32, 2**31 - 1), (np.int64, 2**63 - 1)]:
+            a = np.array(maxv, dtype=t)
+            b = np.array(-2, dtype=t)
+            q, r = mx.divmod(mx.array(a), mx.array(b))
+            np_q, np_r = np.divmod(a, b)
+            self.assertEqual(np.array(q).item(), np_q.item())
+            self.assertEqual(np.array(r).item(), np_r.item())
 
     def test_tile(self):
         self.assertCmpNumpy([(2,), [2]], mx.tile, np.tile)

tests/ops_tests.cpp

@@ -3572,6 +3572,38 @@ TEST_CASE("test divmod") {
   CHECK(array_equal(out[0], array({2.0, 3.0, 3.0})).item<bool>());
   CHECK(array_equal(out[1], array({1.0, 0.0, 1.0})).item<bool>());
 
+  // Negative operands: the quotient must floor toward -inf and the remainder
+  // must carry the divisor's sign, matching numpy (q * b + r == a).
+  x = array({-7, 7, -7, 7});
+  y = array({2, -2, -2, 2});
+  out = divmod(x, y);
+  CHECK(array_equal(out[0], array({-4, -4, 3, 3})).item<bool>());
+  CHECK(array_equal(out[1], array({1, -1, -1, 1})).item<bool>());
+
+  x = array({-7.0, 7.0, -7.0, 7.0});
+  y = array({2.0, -2.0, -2.0, 2.0});
+  out = divmod(x, y);
+  CHECK(array_equal(out[0], array({-4.0, -4.0, 3.0, 3.0})).item<bool>());
+  CHECK(array_equal(out[1], array({1.0, -1.0, -1.0, 1.0})).item<bool>());
+
+  // Signed integer extremes: the naive (a - r) construction overflows here
+  // (e.g. INT_MAX - (-1)), so the quotient must come from shifting the
+  // truncating result instead.
+  x = array({2147483647, -2147483647}, int32);
+  y = array({-2, 2}, int32);
+  out = divmod(x, y);
+  CHECK(array_equal(out[0], array({-1073741824, -1073741824}, int32)).item<bool>());
+  CHECK(array_equal(out[1], array({-1, 1}, int32)).item<bool>());
+
+  x = array({int64_t(9223372036854775807LL), int64_t(-9223372036854775807LL)}, int64);
+  y = array({int64_t(-2), int64_t(2)}, int64);
+  out = divmod(x, y);
+  CHECK(
+      array_equal(out[0], array({int64_t(-4611686018427387904LL),
+                                 int64_t(-4611686018427387904LL)}, int64))
+          .item<bool>());
+  CHECK(array_equal(out[1], array({int64_t(-1), int64_t(1)}, int64)).item<bool>());
+
   x = array({1.0}, complex64);
   y = array({2.0}, complex64);
   CHECK_THROWS(divmod(x, y));