Floor integer floor_divide to match python and numpy - #4346
Open
ayaangazali wants to merge 1 commit into
Open
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follow up to #4108, split out at your request there.
Integer division truncates toward zero, so
mx.floor_dividedisagrees with python and numpy whenever the operands have opposite signs:mx.remainderalready floors, so today//and%describe different divisions, and thedivmoddocstring'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)) / bis that the shorter form overflows: forint8witha = 120andb = -27the numerator is 135, which wraps to -121, and-121 / -27is 4 instead of -5.quotient * bhas the sign ofaand 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 / -1which 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 throughfloor(a / b)untouched, andvmapandcompileof//both match numpy, which matters becauseDivide::vmaprewrites integer division tofloor_divide.test_ops.py,test_vmap.py,test_compile.py,test_autograd.pyandtest_array.pypass, 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.