Skip to content

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

Closed
JemmaUZH wants to merge 1 commit into
ml-explore:mainfrom
JemmaUZH:fix-divmod-floor-semantics
Closed

JemmaUZH wants to merge 1 commit into
ml-explore:mainfrom
JemmaUZH:fix-divmod-floor-semantics

Conversation

@JemmaUZH

@JemmaUZH JemmaUZH commented Aug 16, 2026 •

Copy link
Copy Markdown

Closes #4119

mx.divmod truncates the quotient toward zero when the operands are negative.
mx.divmod(-7, 2) gives you (-3, 1), but numpy gives (-4, 1). And since
(-3) * 2 + 1 isn't -7, the old behavior doesn't even satisfy
q * b + r == a. This shows up for both integer and float operands, on all
three backends.

Here's the repro:

import mlx.core as mx

q, r = mx.divmod(mx.array(-7.0), mx.array(2.0))
print(q.item(), r.item())  # was (-3.0, 1.0); now (-4.0, 1.0)
# numpy reference: np.divmod(-7.0, 2.0) == (-4.0, 1.0)

What I changed

For integers, I start from the truncating quotient and remainder and shift both
down by one when the signs differ. The obvious alternative is to compute
(a - r), but that overflows signed integers at the extremes: INT_MAX / -2
used to come back with the wrong sign because of exactly this.

For floats, the quotient is just floor(a / b), plus numpy's special cases:
when b == 0 the quotient is a / b, and for nan inputs or an infinite
dividend it's nan (a finite dividend that overflows still gives ±inf). It
matters that this computes floor(a / b) directly rather than deriving the
quotient from the remainder; the latter can be off by one ulp.

Metal and CUDA now do the same thing as the CPU backend. The truncating
FloorDivide kernel helper only existed for DivMod, so it's gone.

On the test side I added negative-operand and signed-extreme cases to both the
C++ and Python divmod tests. The Python test now checks the remainder as well,
not just the quotient.

Two things worth flagging

This changes what you get back for negative operands, so anything relying on
the old truncating semantics will see different numbers. Might want a changelog
entry or a minor version bump (your call).

Integer floor_divide has the same root cause but I left it alone.
mx.floor_divide(-7, 2) still returns -3 where numpy says -4, because the
integer path goes through the Divide primitive, so it's a different code path
with its own tests. I'd suggest a follow-up issue, but I'm happy to fold it
into this PR if you'd rather have both fixed at once.

How I verified it (CPU-only build)

The C++ suite passes with the new divmod cases, and so does Python's
test_divmod.

For integers I checked int8/16/32/64 and uint8/16/32/64 across every sign
combination, plus the extremes such as INT64_MAX / -2, INT64_MAX / -3,
INT32_MAX / -2, UINT64_MAX / 3. All match numpy exactly, and
q * b + r == a holds everywhere.

For floats I ran 50k random float16 values and 100k each of float32 and
float64. Every one matches numpy bit for bit, and reconstruction behaves the
same as numpy does, including how it rounds at extreme magnitudes. The float64
numbers are from the CPU backend, since MLX's GPU float64 support is limited.
Edge cases (x/0, ±inf, nan, 0/0, overflow) match numpy too.

I also checked INT_MIN / -1 locally on Apple Silicon. It produces numpy's
(-2147483648, 0) and doesn't trap. I deliberately didn't add it as a test
case though, since it's UB in C++ and would trap on other architectures (x86
idiv, for instance).

I haven't measured performance. The extra work is a sign comparison and a
conditional per element on the integer path, and floor instead of trunc on
the float path. All of it is scalar and only affects the divmod path. Can run
benchmarks if you want numbers.

mx.divmod truncated the quotient toward zero for negative operands, so it
disagreed with mx.floor_divide and with numpy, and broke the q * b + r == a
reconstruction invariant.

The remainder now follows numpy's divisor-sign convention and the quotient
floors like numpy's floor_divide:
- CPU: the integer path shifts the truncating quotient/remainder down by one
  when signs differ (avoiding (a - r), which overflows at signed extremes);
  the float path computes floor(a / b) with numpy's b == 0 and inf/nan
  handling
- Metal/CUDA: same semantics; the truncating FloorDivide helper (only used by
  DivMod) is removed
- Tests: negative operands, remainder checks, and signed-extreme cases added

Verified locally (CPU build): C++ suite passes; Python test passes; random
integer/float inputs match numpy bit for bit, including inf/nan/div-by-zero
edges.

Fixes ml-explore#4119
@ayaangazali

Copy link
Copy Markdown
Contributor

Heads up so you do not spend more time on this than you need to: #4108 is the same fix and has been approved since the 14th. It touches the same three files you do, backend/cpu/binary.cpp, backend/cuda/device/binary_ops.cuh and backend/metal/kernels/binary_ops.h.

The one place they differ is scope. #4108 also floors integer floor_divide in mlx/ops.cpp, because the divmod docstring promises it is equivalent to (a // b, a % b) and that is only true once // floors too. That came out of the review thread, where the ask was to fix all of it rather than just the GPU half.

I am not claiming any priority here, and I genuinely do not mind which one lands. I am only flagging it so the same diff does not get reviewed twice, and so you know the integer floor_divide half exists if whichever PR survives needs to cover it.

One thing from mine that may save you time either way: routing floor_divide through the two output DivMod primitive overruns vmap_replace's out axes vector, which aborts tests/vmap_tests.cpp under ASAN, and Divide::vmap rewrites integer division to floor_divide so a correction applied there lands twice. Worth a check against the C++ suite if your version goes near that path.

@zcbenz

zcbenz commented Aug 16, 2026

Copy link
Copy Markdown
Member

Closing as duplicate of #4108.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] mx.divmod for floats truncates the quotient — disagrees with mx.floor_divide and breaks q·b+r == a

3 participants