Repository navigation
Conversation
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
|
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, The one place they differ is scope. #4108 also floors integer 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 One thing from mine that may save you time either way: routing |
|
Closing as duplicate of #4108. |
Closes #4119
mx.divmodtruncates 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 + 1isn't-7, the old behavior doesn't even satisfyq * b + r == a. This shows up for both integer and float operands, on allthree backends.
Here's the repro:
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 / -2used 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 == 0the quotient isa / b, and for nan inputs or an infinitedividend it's nan (a finite dividend that overflows still gives ±inf). It
matters that this computes
floor(a / b)directly rather than deriving thequotient 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
FloorDividekernel helper only existed forDivMod, 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_dividehas the same root cause but I left it alone.mx.floor_divide(-7, 2)still returns-3where numpy says-4, because theinteger path goes through the
Divideprimitive, so it's a different code pathwith 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, andq * b + r == aholds 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 / -1locally on Apple Silicon. It produces numpy's(-2147483648, 0)and doesn't trap. I deliberately didn't add it as a testcase 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
floorinstead oftrunconthe float path. All of it is scalar and only affects the divmod path. Can run
benchmarks if you want numbers.