Floor integer floor_divide to match python and numpy

mlx/ops.cpp

@@ -3154,10 +3154,23 @@ array floor_divide(
     return floor(divide(a, b, s), s);
   }
 
+  // Integer division truncates toward zero, so take that quotient and step it
+  // down when the result was negative and did not divide evenly. Deriving the
+  // quotient from a - remainder(a, b) instead would leave the dtype range: for
+  // int8 that numerator is 135 when a is 120 and b is -27.
   auto inputs = broadcast_arrays({astype(a, dtype, s), astype(b, dtype, s)}, s);
   auto shape = inputs[0].shape();
-  return array(
-      shape, dtype, std::make_shared<Divide>(to_stream(s)), std::move(inputs));
+  auto quotient =
+      array(shape, dtype, std::make_shared<Divide>(to_stream(s)), inputs);
+  // quotient * b has the sign of a and is no larger in magnitude, so this
+  // truncated remainder is exact for every input the division itself accepts.
+  auto zero = array(0, dtype);
+  auto rem = subtract(inputs[0], multiply(quotient, inputs[1], s), s);
+  auto step = logical_and(
+      not_equal(rem, zero, s),
+      not_equal(less(rem, zero, s), less(inputs[1], zero, s), s),
+      s);
+  return subtract(quotient, astype(step, dtype, s), s);
 }
 
 array remainder(const array& a, const array& b, StreamOrDevice s /* = {} */) {

python/tests/test_ops.py

@@ -330,6 +330,50 @@ def test_multiply(self):
         self.assertEqual(z.dtype, mx.float32)
         self.assertEqual(z.item(), 6.0)
 
+    def test_floor_divide_integers_floor(self):
+        # Integer // truncated toward zero, so it disagreed with python and
+        # numpy whenever the operands had opposite signs.
+        av = [-7, 7, -7, 7, -1, 1, -5, 5, 6, -6]
+        bv = [2, 2, -2, -2, 3, -3, 3, -3, 3, 3]
+        got = mx.floor_divide(mx.array(av), mx.array(bv))
+        self.assertEqual(got.tolist(), [x // y for x, y in zip(av, bv)])
+
+        # The quotient comes from the truncating one plus a correction rather
+        # than from a - remainder(a, b), which can leave the dtype range: for
+        # int8 that numerator is 135 when a is 120 and b is -27.
+        for np_dtype in (np.int8, np.int16):
+            info = np.iinfo(np_dtype)
+            values = [info.min, info.min + 1, -1, 0, 1, 2, info.max - 1, info.max]
+            pairs = [
+                (x, y)
+                for x in values
+                for y in values
+                if y != 0 and not (x == info.min and y == -1)
+            ]
+            a_np = np.array([p[0] for p in pairs], np_dtype)
+            b_np = np.array([p[1] for p in pairs], np_dtype)
+            want = np.floor_divide(a_np.astype(np.int64), b_np.astype(np.int64)).astype(
+                np_dtype
+            )
+            got = mx.floor_divide(mx.array(a_np), mx.array(b_np))
+            self.assertEqual(got.tolist(), want.tolist(), msg=str(np_dtype))
+
+        # Unsigned division already floors and is unchanged.
+        au = np.array([0, 1, 7, 255], np.uint8)
+        bu = np.array([1, 2, 3, 7], np.uint8)
+        self.assertEqual(
+            mx.floor_divide(mx.array(au), mx.array(bu)).tolist(),
+            np.floor_divide(au, bu).tolist(),
+        )
+
+        # Floats keep going through floor(a / b).
+        af = np.array([7.5, -7.5, 7.5, -7.5], np.float32)
+        bf = np.array([2.0, 2.0, -2.0, -2.0], np.float32)
+        self.assertEqual(
+            mx.floor_divide(mx.array(af), mx.array(bf)).tolist(),
+            np.floor_divide(af, bf).tolist(),
+        )
+
     def test_divide(self):
         x = mx.array(2.0)
         y = mx.array(4.0)