Skip to content

Fix floor division with infinite operands - #4335

Open
tandede wants to merge 1 commit into
ml-explore:mainfrom
tandede:agent/fix-floor-divide-infinities
Open

Fix floor division with infinite operands#4335
tandede wants to merge 1 commit into
ml-explore:mainfrom
tandede:agent/fix-floor-divide-infinities

Conversation

@tandede

@tandede tandede commented Aug 18, 2026

Copy link
Copy Markdown

Summary

  • add a dedicated single-output FloorDivide primitive for floating-point inputs on CPU, Metal, and CUDA
  • match Python and NumPy for infinite operands, including finite values divided by oppositely signed infinities and infinite numerators
  • keep the quotient returned by divmod consistent with floor_divide for the same edge cases
  • preserve the integer path and explicitly reject unsupported complex inputs
  • support compilation, vmap, and autodiff for the new primitive, with regression coverage for eager, compiled, and divmod execution

Root 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, while floor(inf / 3) remains infinite instead of becoming NaN.

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 passed
  • CPU C++ suite: 251 test cases and 3,358 assertions passed
  • python/tests/test_ops.py: 162 passed, 1 skipped
  • all 49 combinations drawn from NaN, +/-Inf, positive and negative finite values, and signed zero matched NumPy values and sign bits
  • full CPU-only Python discovery: 801 passed and 40 skipped; the remaining 18 MPS DLPack cases require Metal and failed because this local build intentionally uses MLX_BUILD_METAL=OFF

On 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

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (no API documentation change is required)

@tandede
tandede marked this pull request as ready for review August 19, 2026 11:28
@tandede
tandede force-pushed the agent/fix-floor-divide-infinities branch from 4e9de8f to 5413e29 Compare August 20, 2026 03:09
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.

[BUG] mx.floor_divide disagrees with NumPy and Python at infinity

2 participants