diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 5bc69eb2f5..355a178c7a 100644 --- a/mlx/ops.cpp +++ b/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(to_stream(s)), std::move(inputs)); + auto quotient = + array(shape, dtype, std::make_shared(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 /* = {} */) { diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 054a1ca6b3..e3418b6a43 100644 --- a/python/tests/test_ops.py +++ b/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)