Fix floor division with infinite operands - #4335
Open
tandede wants to merge 1 commit into
Open
Conversation
tandede
marked this pull request as ready for review
August 19, 2026 11:28
tandede
force-pushed
the
agent/fix-floor-divide-infinities
branch
from
August 20, 2026 03:09
4e9de8f to
5413e29
Compare
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.
Summary
FloorDivideprimitive for floating-point inputs on CPU, Metal, and CUDAdivmodconsistent withfloor_dividefor the same edge casesvmap, and autodiff for the new primitive, with regression coverage for eager, compiled, anddivmodexecutionRoot cause
Floating-point floor division was implemented as
floor(a / b). That is not equivalent to Python floor division at infinity:floor(1 / -inf)preserves negative zero instead of returning-1, whilefloor(inf / 3)remains infinite instead of becomingNaN.The fix uses a dedicated primitive so the exceptional cases can be corrected in one backend pass without computing an unused remainder or materializing an intermediate quotient array. The CPU implementation retains a SIMD fast path when neither input contains infinities.
Closes #4317.
Validation
pre-commit run --all-files: all hooks passedpython/tests/test_ops.py: 162 passed, 1 skippedNaN,+/-Inf, positive and negative finite values, and signed zero matched NumPy values and sign bitsMLX_BUILD_METAL=OFFOn one million finite float32 elements, the dedicated CPU primitive reduced the local median evaluation time from 0.248 ms to 0.202 ms across seven 20-iteration samples (about 1.23x). Metal and CUDA compilation and runtime coverage are left to CI because those toolchains are not available locally.
Checklist
pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes