Derive the divmod remainder from its quotient on the GPU

[ayaangazali]

Proposed changes

Fix #4119.

DivMod on the GPU pairs a truncating quotient with a flooring remainder, so the two do not describe the same division:

// mlx/backend/metal/kernels/binary_ops.h and cuda/device/binary_ops.cuh
return {FloorDivide{}(x, y), Remainder{}(x, y)};

FloorDivide truncates despite the name (x / y for integers, trunc(x / y) for floats). Remainder floors, since it applies the usual sign correction:

auto r = x % y;
if (r != 0 && (r < 0 != y < 0)) {
  r += y;
}

Mixing them breaks quotient * y + remainder == x. Working the two structs through by hand for int32, on the eight mixed-sign pairs below, six come out wrong:

   a    b   quotient  remainder   q*b+r
  -7    2      -3         1        -5    should be -7
   7   -2      -3        -1         5    should be  7
  -1    3       0         2         2    should be -1
   1   -3       0        -2        -2    should be  1
  -5    3      -1         1        -2    should be -5
   5   -3      -1        -1         2    should be  5

The CPU backend does not have this problem, because it builds both halves from the same division:

// mlx/backend/cpu/binary.cpp
return std::make_pair(x / y, x % y);

and mx.divmod there satisfies the identity on all eight pairs. So this is the GPU disagreeing with the CPU as well as with itself.

Taking the remainder from the quotient that is actually returned, which is the same thing the CPU does and leaves the quotient untouched:

auto quotient = FloorDivide{}(x, y);
return {quotient, x - quotient * y};

test_divmod did not catch it because it draws inputs from np.random.uniform(1, 100), so both operands are always positive, and it only compares mx_out[0], never the remainder. The added case checks the identity on mixed signs.

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)

Being explicit about what I could and could not run: my box is a CPU only build, so I could not execute the Metal or CUDA path. The table above comes from working the two structs through by hand rather than from a GPU run, and test_ops.py (149 tests) passes here. The new assertion passes on CPU both before and after, since the CPU was already consistent, so it is the GPU jobs in CI that should show it red on main and green here. If they do not, my reading of the two structs is wrong and I would rather find that out than have it merged.

---

freshman contributor, i work through these with Claude Code. i got here sideways: i was diffing the same op across the three backends looking for places where one had a guard the others lacked, noticed FloorDivide exists only on the GPU side, and followed it to its one caller.

[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.