Fix divmod truncating the quotient for floats

[ayaangazali]

Fix #4119

Proposed changes

mx.divmod truncated its quotient while mx.remainder floored, so the two disagreed with each other and with everything else:

>>> mx.divmod(mx.array([-7]), mx.array([2]))     # was (-3, -1)
>>> mx.remainder(mx.array([-7]), mx.array([2]))  # 1
>>> divmod(-7, 2)                                # (-4, 1)

numpy, pytorch and python all floor both halves. The docstring for divmod says it is "equivalent to but faster than (a // b, a % b)", which was not true: a // b truncated for integers and a % b floored, so that pair did not even satisfy q * b + r == a.

On the GPU it was worse, because DivMod paired a truncating FloorDivide with a flooring Remainder, so quotient * y + remainder == x failed outright on six of eight mixed-sign int32 pairs.

Everything now floors:

• FloorDivide on Metal and CUDA steps the quotient down when the operands have opposite signs and the division is not exact, and uses floor rather than trunc for floats. DivMod goes back to {FloorDivide, Remainder}, which is consistent now that both halves floor.
• DivMod::eval_cpu applies the same correction, and its float path returns (floor(x / y), floored remainder).
• Integer floor_divide takes the floored quotient from divmod instead of the truncating Divide primitive, so the two cannot drift apart again.

>>> a = mx.array([-7, 7, -7, 7, -1, 1, -5, 5])
>>> b = mx.array([2, 2, -2, -2, 3, -3, 3, -3])
>>> mx.divmod(a, b)[0].tolist()
[-4, 3, 3, -4, -1, -1, -2, -2]
>>> [x // y for x, y in zip(a.tolist(), b.tolist())]
[-4, 3, 3, -4, -1, -1, -2, -2]

Quotient and remainder both match python and numpy on every mixed-sign pair I tried, q * b + r == a holds, and a // b and a % b now equal the two halves of divmod. Floats match divmod too. Unsigned types are untouched, since the correction is behind is_signed_v, and dtypes are preserved.

One tradeoff worth calling out: routing integer floor_divide through divmod computes a remainder that is then dropped. I preferred that over open coding the correction a second time, since the whole point here is that the two stopped agreeing, but say the word and I will inline it instead.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

CPU only build (MLX_BUILD_METAL=OFF) at 3abd0fd. test_ops.py, test_array.py, test_autograd.py, test_compile.py, test_random.py, test_losses.py, test_nn.py and test_linalg.py all pass, and the new assertions fail on main. I cannot run Metal or CUDA here, so those two kernels are reasoned rather than executed and the GPU jobs are the real check on them.

Also worth noting: mx.divmod(-6.0, 3.0) gives a remainder of -0.0 where numpy gives 0.0. That is unchanged behaviour and it now matches mx.remainder, which is what pytorch does too.

---

freshman contributor, i work through these with Claude Code. thanks for the steer on option 2, this ended up much tidier than what i had. the part i had to be careful with was unsigned integers, where the sign correction has to be compiled out rather than just evaluated to false, otherwise the always false comparison trips the warnings as errors build.

[zcbenz]

Can you check how numpy/pytorch behaves on this?

[ayaangazali]

Checked, and the answer argues for a bigger change than this PR makes. Measured on int32, mlx via the CPU stream:

   a,  b | np.divmod | torch floor_divide,remainder | python divmod | mx.divmod
  -7,  2 |   (-4, 1) |                      (-4, 1) |       (-4, 1) |  (-3, -1)
   7, -2 |  (-4, -1) |                     (-4, -1) |      (-4, -1) |   (-3, 1)
  -1,  3 |   (-1, 2) |                      (-1, 2) |       (-1, 2) |   (0, -1)
   1, -3 |  (-1, -2) |                     (-1, -2) |      (-1, -2) |    (0, 1)
  -5,  3 |   (-2, 1) |                      (-2, 1) |       (-2, 1) |  (-1, -2)
   5, -3 |  (-2, -1) |                     (-2, -1) |      (-2, -1) |   (-1, 2)

numpy, pytorch and python all floor both halves. mlx truncates both. All four satisfy q*b + r == a, so this is a convention difference rather than a correctness one, and pytorch does offer a truncating pair too (div(rounding_mode='trunc') with fmod), which is also self consistent.

So there are two ways to go:

1. What this PR does. Keep truncation and make the GPU match the CPU. Smallest change, no behaviour change for anyone already on CPU, but mlx keeps disagreeing with numpy on mixed signs.
2. Floor both halves everywhere. That matches numpy, pytorch and python, and it also lines up with mx.remainder, which already floors, so mx.divmod(a, b)[1] and mx.remainder(a, b) would stop disagreeing. It changes CPU results for mixed signs, and to make the docstring's "equivalent to (a // b, a % b)" true it would also need integer floor_divide to floor, which today truncates.

I did not want to pick option 2 for you inside a PR scoped to the GPU, so tell me which you want and I will redo it that way. Also pushed a fix for the metal build failure, which was mine: for the narrow int types x - quotient * y widened to int and the braced initializer would not narrow it back.

[zcbenz]

Thanks for checking, let's do option 2 and fix all things.

[ayaangazali]

Nudging this one with a recommendation rather than just a ping, since the measurement came back one sided.

Every reference floors both halves: numpy, pytorch and python divmod all return (-4, 1) for divmod(-7, 2), and mlx returns (-3, -1). mlx's own mx.remainder already floors and matches them, so today mx.divmod(a, b)[1] and mx.remainder(a, b) disagree with each other.

If you want option 2, I am happy to redo this as: floor both halves in DivMod on all three backends, and make integer floor_divide floor so the docstring's "equivalent to (a // b, a % b)" becomes true. That changes CPU results for mixed signs, which is why I did not just do it.

If you would rather keep truncation, this PR as it stands is the smaller change and only makes the GPU agree with the CPU. Either way is fine by me, I just did not want to pick the convention on your behalf. Happy to close this if you prefer the bigger change as a fresh PR.

[ayaangazali]

Ignore that last comment, I missed your reply above. Doing option 2.

[zcbenz]

This seems to be causing a lot of weird errors, adding a new primitive might be simpler in the end.

[ayaangazali]

Pushed fixes for both CI failures. Both were mine.

Metal build. My FloorDivide rewrite kept the old template <> specializations for float, half and bfloat16 while constraining the primary on !is_signed_v<T>, which those three do not satisfy, so the specializations had nothing to attach to:

binary_ops.h:34:9: error: no function template matches function template specialization 'operator()'

It now uses the same three enable_if_t overloads that Remainder above it uses (unsigned integral, signed integral, non integral), with floor in the non integral one, so no explicit specializations are needed.

Fedora and ASAN crash. Routing integer floor_divide through divmod was simply wrong: DivMod has two outputs, so vmap_replace read past the end of the out axes vector and tests/vmap_tests.cpp:561 aborted.

Worse, that path was silently wrong under vmap even where it did not crash. Divide::vmap rewrites integer division to floor_divide, so my truncating quotient came back already floored and the correction applied twice: vmap of floor_divide(-12, 5) gave -4 instead of -3.

So floor_divide no longer touches the multi output primitive. It subtracts the floored remainder first, which makes the division exact, so it does not matter which way the division rounds and the vmap rewrite is harmless:

auto num = subtract(a, remainder(a, b, s), s);
// ... Divide primitive on (num, b)

That also drops the wasted remainder computation you would have had, so the tradeoff I flagged earlier is gone.

Verified: all 216 combinations of a in [-13, 13] and b in {1, 2, 3, 5, -1, -2, -3, -5} match python for divmod, // and %; vmap matches python for divisors 5, -5, 3 and -3 across the same range; unsigned and every integer width keep their dtype; floats unchanged. I also built and ran the C++ suite this time rather than only the python one, 249 cases and 3346 assertions pass, including the test that was crashing.

[ayaangazali]

Third metal fix, sorry for the churn. The build step failed again and it was the complex instantiation.

binary_two.metal does instantiate_binary_types(DivMod), and that list includes complex64, so FloorDivide has to compile for complex64_t even though divmod rejects complex at the op level long before a kernel runs. My previous version routed complex into the non integral overload, which calls floor, and there is no floor for a complex.

FloorDivide now has the same four overloads Remainder directly above it has: unsigned integral, signed integral, non integral, and a template <> for complex64_t. I added the matching is_complex_v branch to the CUDA one too, since its Remainder has one and the kernels there are instantiated the same way.

Ran the C++ suite again as well as the python one, 249 cases and 3346 assertions pass. I still cannot compile Metal or CUDA here, so the structure being identical to Remainder is the strongest check I can make locally.

[zcbenz]

Thanks!

[ayaangazali]

Re-verified this on current main (bbebc8f), since it has been sitting a week and main has moved a lot in that time.

It still merges cleanly, and on that merge the C++ suite passes (250 cases, 3351 assertions) along with test_ops.py, test_autograd.py, test_vmap.py, test_compile.py and test_array.py. Behaviour is unchanged from what you approved:

   a, b | numpy   | python  | mx.divmod | mx.//  | mx.%
  -7,  2 | (-4, 1) | (-4, 1) |  (-4, 1)  |  -4    |  1
   7, -2 | (-4,-1) | (-4,-1) |  (-4,-1)  |  -4    | -1
  -1,  3 | (-1, 2) | (-1, 2) |  (-1, 2)  |  -1    |  2
   5, -3 | (-2,-1) | (-2,-1) |  (-2,-1)  |  -2    | -1

q * b + r == a holds for every sign combination.

I have deliberately not force pushed a rebase. Fork PRs are behind the workflow approval gate at the moment, so a push would drop the 28 green checks here back to zero and need another approval to re-run. Nothing needs rebasing anyway, GitHub still reports it mergeable and clean.

One thing worth knowing rather than discovering twice: #4311 was opened today and covers the same three backend files. I have left a note there too. I do not mind at all which one you take, I just did not want the same diff reviewed twice.

[krrishapatel]

Two things I hit running your branch merged onto main.

1. The integer path in mlx/ops.cpp overflows. a - remainder(a, b) can leave the dtype range:

a, b = mx.array([120, 1], mx.int8), mx.array([-27, -128], mx.int8)

(a // b).tolist()              # [4, 1]
mx.divmod(a, b)[0].tolist()    # [-5, -1]   correct

remainder(120, -27) is -15, so a - r is 135, which wraps to -121 in int8, and -121 / -27 is 4. So // and divmod disagree, which is the invariant this PR is for. int16 too. int32 and int64 need operands near the boundary.

2. floor(x / y) divides a second time, so the quotient can disagree with the remainder:

a, b = mx.array([2144.0], mx.bfloat16), mx.array([358.0], mx.bfloat16)
mx.divmod(a, b)     # q=6.0, r=354.0    2144/358 is 5.9888, so the floor is 5

5.9888 rounds to exactly 6.0 in bfloat16 and floor cannot undo it. In float32 divmod(1.0, 0.1) gives q=10 where numpy gives 9. The new float test passes because 7.5 / 2.0 is exact.

numpy derives the quotient from the remainder instead of dividing twice:

mod = fmod(a, b);
div = (a - mod) / b;
if (mod != 0 && (b < 0) != (mod < 0)) { mod += b; div -= 1; }
else if (mod == 0) { mod = copysign(0, b); }
floordiv = floor(div);
if (div - floordiv > 0.5) floordiv += 1;

I checked that against np.divmod over all 400 pairs of a 20 value set in float16, float32 and float64. No mismatches, where floor(x / y) misses 16, 34 and 30.

Found with https://github.com/krrishapatel/arraydiff.

[krrishapatel]

Correction to the algorithm I posted above. I implemented it in binary.cpp to check it, and my listing was missing two branches numpy has.

A zero divisor needs an early out. Otherwise fmod is NaN and (a - mod) / b carries that into the quotient. numpy and mlx main both give the signed infinity there, so without this it is a regression.

A zero quotient needs its own sign. (a - mod) / b has already cancelled to zero, so it carries the sign of the cancellation rather than of the division. divmod(-0.1, -1.0) should be +0.

mod = fmod(a, b);
if (b == 0) { return {a / b, mod}; }                 // was missing
div = (a - mod) / b;
if (mod != 0) { if ((b < 0) != (mod < 0)) { mod += b; div -= 1; } }
else { mod = copysign(0, b); }
if (div != 0) {
  floordiv = floor(div);
  if (div - floordiv > 0.5) { floordiv += 1; }
} else {
  floordiv = copysign(0, a / b);                     // was missing
}

With both branches this is bit exact against np.divmod over 625 pairs in float16, float32 and float64, signed zeros and inf and nan included. The two I left out were 45 and 24 mismatches.

[krrishapatel]

I should have checked #4003 before posting the algorithm above. It found the same half precision rounding and was closed with "we should follow what PyTorch does" and prefer speed. That is worth revisiting here, because PyTorch's floor division is not floor(a / b).

c10/util/generic_math.h:

if (C10_UNLIKELY(b == 0)) { return a / b; }
auto mod = std::fmod(a, b);
auto div = (a - mod) / b;
if ((mod != 0) && (b < 0) != (mod < 0)) { div -= scalar_t(1); }
scalar_t floordiv;
if (div != 0) {
  floordiv = std::floor(div);
  if (div - floordiv > scalar_t(0.5)) { floordiv += scalar_t(1.0); }
} else {
  floordiv = C10_COMPAT_COPYSIGN(scalar_t(0), a / b);
}

That derives the quotient from the remainder, same as numpy, and it is vectorized as div_floor_floating_vec in BinaryOpsKernel.cpp. The torch.div example in that review is plain division, which is a different path.

It matters for this PR because of the new comment, "Both halves floor, so quotient * y + remainder == x". I applied this PR's float_op on CPU and that does not hold:

mx.divmod(mx.array([2144.0], mx.bfloat16), mx.array([358.0], mx.bfloat16))
# q=6, r=354      q*y + r = 2496, x = 2144

mx.divmod(mx.array([1.0], mx.float64), mx.array([0.1], mx.float64))
# q=10, r=0.09999999999999995      q*y + r = 1.0999999999999999, x = 1.0

float64 is the part #4003 does not reach. Its fix was to widen the division, and there is nothing above float64 to widen to, so "use higher precision" has no answer there. The error is also a whole divisor rather than a rounding difference.

I have the PyTorch version building on CPU: bit exact against np.divmod over 625 pairs per precision in float16, float32 and float64, signed zeros and inf and nan included, and q*y + r == x in all four including bfloat16. Glad to send it as a follow up here, or to drop it if you would rather keep the branches off the hot path.

[ayaangazali]

@krrishapatel thank you, the first one is a real regression in this PR and it should not be merged as it stands. Please hold off on this one for now.

I reproduced both on this branch merged onto current main, and I want to separate them because they are not the same kind of problem.

The integer overflow is mine, and it is worse than the invariant it was meant to fix. On main // and divmod agree with each other, both truncating. With this PR divmod becomes right and // becomes wrong:

             a // b      divmod q     numpy
main         [-4,  0]    [-4,  0]     [-5, -1]
this PR      [ 4,  1]    [-5, -1]     [-5, -1]

for a = int8([120, 1]), b = int8([-27, -128]). So the PR trades a convention disagreement for a wrong answer, and it breaks agreement between // and divmod, which is the whole point of the change. Your diagnosis is exactly right: a - remainder(a, b) is 135 for the first pair, which does not fit in int8.

The float one reproduces but predates this PR. Same numbers on unmodified main:

bfloat16 divmod(2144, 358)   main: q=6 r=354   this PR: q=6 r=354
float32  divmod(1.0, 0.1)    main: q=10        this PR: q=10     numpy: 9

So this branch does not change that behaviour. What it does do is add a comment asserting "Both halves floor, so quotient * y + remainder == x", and your examples show that sentence is not true for the float path. The claim is wrong even though the code is not new, and I will fix the claim.

I am working on the integer fix now: derive the floored quotient from the truncating one and a correction term rather than from a - remainder, so nothing has to leave the dtype range. I will post the result here.

On your PyTorch and numpy algorithm for the float path, I agree that deriving the quotient from the remainder is the correct answer, and your point about float64 having nothing to widen to is a good one against the #4003 resolution. That is a bigger change than this PR and it is your work, so please do send it as its own PR rather than folding it in here. I will point at it from here.

[ayaangazali]

Fixed and pushed.

The floored quotient now comes from the truncating one plus a correction, so nothing has to leave the dtype range:

auto quotient = array(shape, dtype, std::make_shared<Divide>(to_stream(s)), inputs);
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);

quotient * b has the sign of a and is no larger in magnitude, so that remainder is exact for every input the division itself accepts.

Your case now agrees:

             a // b      divmod q     numpy
before       [ 4,  1]    [-5, -1]     [-5, -1]
now          [-5, -1]    [-5, -1]     [-5, -1]

Verified over all 65279 int8 pairs with a non zero divisor, skipping only INT_MIN / -1 which overflows before any of this:

• floor_divide matches numpy on every pair, 0 mismatches
• remainder matches numpy on every pair
• divmod agrees with floor_divide and remainder on every pair
• q * b + r == a on every pair

Also checked the boundary values for int16, int32, int64 and all four unsigned types, and vmap, nested vmap and compile of // against a plain loop, since routing this through the multi output primitive was what broke vmap_replace earlier. The float path is untouched.

C++ suite passes (249 cases, 3346 assertions) along with test_ops.py, test_autograd.py, test_vmap.py, test_compile.py, test_array.py and test_reduce.py. Added a regression test over the boundary values for int8 and int16, which fails on the previous commit.

I also corrected the comment you were right about. It now says the identity holds for integers and explicitly does not claim it for floats, since flooring an already rounded quotient cannot recover the exact one.

Thanks for catching this, it was a clear regression and it had my approval on it. Your float work stands on its own and I would rather it went in as your PR than folded in here.

[zcbenz]

I think with this implementation we should really add a new primitive for FloorDivide, probably in another PR.

[ayaangazali]

Agreed, and I think you are right that a FloorDivide primitive is where this wants to end up.

What this PR does now is build the floored quotient out of existing primitives, a Divide plus a multiply, a subtract, two compares and a select. That is correct, and it is verified over all 65279 int8 pairs, but it is six nodes in the graph where one kernel would do, and it is why the overflow was possible in the first place: the intermediate had to fit in the dtype. A primitive computes the quotient and the correction in registers and never materializes anything that can wrap.

Happy to do it as a follow up rather than growing this PR, since a new primitive means the three backend kernels, vjp and vmap, and it should land on its own. If you would rather have that instead of this and skip the intermediate step, I am fine closing this one too, your call.

[zcbenz]

I think it would be good to have this PR fixing divmod, and then a followup PR to fix floor_divide.

[ayaangazali]

Done, this PR is now divmod only.

mlx/ops.cpp is back to what main has, so floor_divide is untouched here, and the tests no longer assert that mx.// agrees with divmod, since it will not until the followup lands. What remains is the three DivMod kernels plus tests that check divmod against python's divmod and the q * b + r == a invariant:

divmod q : [-4, 3, 3, -4, -1, -1, -2, -2]
python   : [-4, 3, 3, -4, -1, -1, -2, -2]
q*b+r==a : True
a // b   : [-3, 3, 3, -3,  0,  0, -1, -1]   still truncating, as expected for now

Floats are unchanged and still satisfy the invariant. C++ suite passes (249 cases, 3346 assertions) with test_ops.py, test_vmap.py, test_compile.py and test_autograd.py.

The followup for integer floor_divide is written and verified against all 65279 int8 pairs plus the boundaries for the wider signed and unsigned types, but I cannot open it right now: GitHub is refusing pull request creation for my account on this repo, through both gh pr create and the REST endpoint. The branch is pushed at ayaangazali:floor-divide-integers if you want to look before I can get the PR up, and I will open it as soon as that clears.