Skip to content

Floor integer floor_divide to match python and numpy - #4346

Open
ayaangazali wants to merge 1 commit into
ml-explore:mainfrom
ayaangazali:floor-divide-integers
Open

Floor integer floor_divide to match python and numpy#4346
ayaangazali wants to merge 1 commit into
ml-explore:mainfrom
ayaangazali:floor-divide-integers

Conversation

@ayaangazali

Copy link
Copy Markdown
Contributor

Follow up to #4108, split out at your request there.

Integer division truncates toward zero, so mx.floor_divide disagrees with python and numpy whenever the operands have opposite signs:

>>> av = [-7, 7, -7, 7, -1, 1, -5, 5]
>>> bv = [ 2, 2, -2, -2, 3, -3, 3, -3]
>>> mx.floor_divide(mx.array(av), mx.array(bv)).tolist()
[-3, 3, 3, -3, 0, 0, -1, -1]
>>> [x // y for x, y in zip(av, bv)]        # python
[-4, 3, 3, -4, -1, -1, -2, -2]
>>> np.floor_divide(av, bv).tolist()        # numpy
[-4, 3, 3, -4, -1, -1, -2, -2]

mx.remainder already floors, so today // and % describe different divisions, and the divmod docstring's claim that it is equivalent to (a // b, a % b) cannot hold for either.

The fix takes the truncating quotient and steps it down when the operands have opposite signs and the division was not exact. The reason it is written that way rather than as the shorter (a - remainder(a, b)) / b is that the shorter form overflows: for int8 with a = 120 and b = -27 the numerator is 135, which wraps to -121, and -121 / -27 is 4 instead of -5. quotient * b has the sign of a and is never larger in magnitude, so the truncated remainder here is exact for every input the division itself accepts.

Verified over all 65279 int8 pairs with a non zero divisor, skipping only INT_MIN / -1 which overflows before any of this, with zero mismatches against numpy. Also the boundary values (min, min+1, -1, 0, 1, 2, max-1, max) for int16, int32, int64 and all four unsigned widths, again zero mismatches. Unsigned division already floored and is unchanged, floats still go through floor(a / b) untouched, and vmap and compile of // both match numpy, which matters because Divide::vmap rewrites integer division to floor_divide.

test_ops.py, test_vmap.py, test_compile.py, test_autograd.py and test_array.py pass, along with the C++ suite (250 cases, 3351 assertions). The added test fails on main.

One consequence worth stating: with this and #4108 both in, mx.divmod(a, b) and (a // b, a % b) agree for integers and the docstring becomes true. Either one alone leaves them disagreeing, in opposite directions.

I am a freshman working through this codebase and I used Claude Code alongside it. Every number above came from a run on this machine.

Integer division truncates toward zero, so mx.floor_divide disagreed
with python and numpy whenever the operands had opposite signs:
-7 // 2 gave -3 where both give -4. Derive the floored quotient from the
truncating one plus a correction, which keeps every intermediate inside
the dtype.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants