fix(ops): floor_divide floors for negative signed integers

[davidtai]

Proposed changes

The defect

floor_divide on integer arrays rounds toward zero. It must round down. -7 // 2 gives -3. NumPy (np.floor_divide) and Python (//) give -4. A test in the Swift wrapper found this.

• mlx/ops.cpp:3142-3155 (floor_divide): for float types it returns floor(divide(a, b)), which is correct. For integer types (and bool) it makes a Divide primitive on the integer inputs.
• mlx/backend/cpu/binary_ops.h:28 on main: DEFAULT_BINARY_OP(Divide, operator/). This is C++ /, which truncates toward zero.
• mlx/backend/metal/kernels/binary_ops.h:43-48 on main: Divide returns x / y. Metal integer / also truncates toward zero.
• The Python __floordiv__, __rfloordiv__ and __ifloordiv__ (python/src/array.cpp:691-719) call mx::floor_divide. They have the same defect.

What happened before this PR:

Type	Before
Signed integers (int8, int16, int32, int64)	Truncated toward zero. Wrong when the signs are different and the remainder is not zero. Correct for all other inputs.
Unsigned integers	Correct. Truncation and floor are the same for values that are not negative.
Floats	Correct: floor(a / b). These do not use the integer Divide path.

Integer Divide exists only for floor_divide. divide promotes integers to float first (mlx/ops.cpp, at_least_float). Divide::vmap already maps an integer Divide to floor_divide (mlx/primitives.cpp:1885). divmod was already correct on all backends: the CPU DivMod code and the Metal and CUDA FloorDivide functors floor the quotient.

Upstream status

Upstream has the same defect up to 2026-09-28. Issue ml-explore#4450 reported it. Upstream fixed it in ml-explore#4515 (commit 0a8e18e0). That fix is in ops.cpp. It keeps the truncating Divide and adds 7 more graph nodes (remainder, two less, two not_equal, logical_and, astype, subtract) to subtract 1. Upstream did not keep truncation on purpose. It calls truncation a bug.

This PR does not port ml-explore#4515 as it is, for two reasons:

1. With Fix floor_divide for integers ml-explore/mlx#4515, division by zero changes on the CPU. Today -7 // 0 gives 0 on the CPU (arm64 sdiv returns 0). With Fix floor_divide for integers ml-explore/mlx#4515, remainder(-7, 0) is -7, which is not 0, and the signs are different, so the result becomes -1.
2. Fix floor_divide for integers ml-explore/mlx#4515 adds 7 kernels to each integer floor_divide. The fix in the functor adds none.

If a later upstream sync ports ml-explore#4515 on top of these flooring kernels, floor_divide subtracts 1 twice (-7 // 2 gives -5). The new tests catch this.

The fix, per backend

For signed integers: q = a / b (truncated). If the remainder is not zero and the signs of a and b are different, then q = q - 1.

• CPU (mlx/backend/cpu/binary_ops.h): Divide is now its own struct. It does the step above for signed integer T, for SIMD widths N > 1 and for N = 1. C++ does not define x / 0 or min / -1, so the functor divides by 1 in those lanes (d = select(y == 0 || (x == min && y == -1), 1, y)). Then it selects 0 where y == 0. It subtracts select(mask, 1, 0), so it does not compute q - 1 in lanes that do not need it (INT_MIN // 1 does not overflow). All other types still use x / y. The CPU compile path (compiled_cpu) uses the same functor.
• Metal (mlx/backend/metal/kernels/binary_ops.h): Divide has two overloads, in the same form as the existing FloorDivide and Remainder. The signed-integer overload does the step above. The other overload (unsigned, bool, float, half, bfloat16, complex) still returns x / y. Fused kernels from mx.compile use the same header.
• CUDA: no change. cu::Divide (mlx/backend/cuda/device/binary_ops.cuh:36-41) still returns x / y. On CUDA, integer floor_divide still truncates. The CUDA file already has a correct FloorDivide functor, which DivMod uses. This PR does not touch CUDA, so the CUDA build does not change. No workflow in this fork builds CUDA.

Division by zero does not change. On main the CPU computes x / 0 with C++ / (DEFAULT_BINARY_OP(Divide, operator/), Simd<T, 1> and the clang vector / for N > 1). C++ does not define this, but arm64 sdiv and udiv return 0, so main gives 0 for signed and unsigned integers on the CPU. The new signed-integer code does not divide by 0: it returns 0 explicitly, so the result does not depend on the compiler. Unsigned integers still use x / y, as on main. On Metal, the Metal Shading Language does not define integer division by zero. The value comes from the GPU hardware and this PR does not change it. For floats, the result is floor(±inf) or NaN, as before. (Note: divmod already gives -1 for -7 divmod 0 on the CPU. This PR does not change divmod.)

The overflow case INT_MIN // -1 does not change. On main the CPU gives INT_MIN (arm64 sdiv wraps; int8 and int16 promote to int and convert back). The new CPU code returns INT_MIN explicitly and does not divide. On Metal this case is not defined and is not tested.

Call-site risk

The behavior changes only for signed integers with operands of different signs and a remainder that is not zero.

This repository

• mlx/primitives.cpp:1885, Divide::vmap: uses floor_divide for integer Divide. It now gives the same result as the unbatched op.
• python/mlx/optimizers/schedulers.py:58: step // step_size. The step count is not negative. No change.
• All other // in python/mlx are on Python ints (shapes, sizes), not arrays.

mlx-swift (Source/, read only, at 0f4fe40)

• Source/MLX/Ops+Array.swift:752 (floorDivide(_:_:stream:)) and Source/MLX/MLXArray+Ops.swift:1913 (MLXArray.floorDivide(_:stream:)) are wrappers around mlx_floor_divide. They pass on the new behavior. The other hits are documentation. A test in the wrapper found this defect, so code in the wrapper that expects -4 now passes.

mlx-swift-lm (Libraries/, read only, at e22fc82b)

Call site	Operands	Negative operands possible?
MLXLLM/Models/Qwen4ExpCompactQSA.swift:35 floorDivide(positions + 1, 4)	positions from int32Range(offset, offset + width), with precondition(offset >= 0)	No
MLXLLM/Models/Qwen4ExpQSA.swift:479 floorDivide(int32Range(queryStart, keyTokens) + 1, ratio)	queryStart = keyTokens - queryTokens (not negative), ratio is the compress ratio (positive)	No
MLXLLM/Models/Qwen4ExpQSA.swift:883 same expression	same	No
MLXLLM/Models/Qwen4ExpQSA.swift:995 floorDivide(tokInt, ratio)	token row indices (not negative), ratio positive	No
MLXLLM/Models/Qwen4ExpQSA.swift:999 floorDivide(pos + 1, ratio)	pos from int32Range(queryStart, keyTokens)	No
MLXLMCommon/SwitchLayers.swift:250 order.floorDivide(m)	order from argSort (uint32), m = indices.dim(-1)	No (unsigned path, no change)

No call site found can get negative operands, so none depends on the old truncation. A related case: the DeepseekV4 pool mask used true division where the reference uses floor division. That is fixed separately and does not use floor_divide.

Downstream: mlx-swift

mlx-swift keeps generated copies of both headers: Source/Cmlx/mlx-generated/binary_ops.cpp, Source/Cmlx/mlx-generated/metal/binary_ops.h and Source/Cmlx/mlx-generated/compiled_preamble.cpp. At the mlx-swift pin bump, tools/update-mlx.sh must run. If it does not, the Swift Metal path still truncates while the CPU floors.

Test plan

• C++ (tests/ops_tests.cpp, test floor_divide with negative operands): on Device::cpu, and on Device::gpu when Metal is available. For int8, int16, int32 and int64: all four sign combinations, exact multiples, zero dividends, and quotients between -1 and 0. Also a scalar divisor, transposed (strided) inputs, a match with divmod(a, b)[0], the int8 limits (-128 // 3 == -43, 127 // -2 == -64), unsigned types (no change) and float32 (no change). A 40-element case for each signed width, so the SIMD path runs (int8 and int16 have 16 lanes). test floor_divide integer edge cases on the CPU (CPU only): floor_divide({-7, 7, 0}, 0) == {0, 0, 0}, a 33-element x // 0 == 0, and a 33-element min // -1 == min for each signed width.
• Python (python/tests/test_ops.py, test_floor_divide_negative_operands): the same cases on the default device, checked against np.floor_divide and Python //, through mx.floor_divide, //, __rfloordiv__ (-7 // y), divmod, transposed inputs and an mx.compile fused kernel (u // v + 1). Also unsigned, float16 and float32. The same 40-element case, and on mx.cpu the same x // 0 and min // -1 cases. Metal does not define these two cases, so they are not tested on the GPU.
• The new tests fail on main: on main, Divide truncates, so -7 // 2 gives -3, and the tests expect -4. Upstream's test for Fix floor_divide for integers ml-explore/mlx#4515 failed in the same way before its fix.
• Fork Tests (metal-build-and-test, macos-26) on head 34b31e90: run 36943701359 passed: 292 C++ test cases with 3945 assertions, and 876 Python tests with 52 skipped. Build and Test run 36943701360 and Fork Diff Check run 36943701354 also passed.
• /usr/bin/python3 scripts/check_forkdiff.py --upstream-ref refs/remotes/upstream/main on the head commit: 173 of 173 changed files described. fork.yaml has a new section, "Integer floor_divide rounds down".
• The commit is signed.

The Python tests use only the default device. The CPU pass in #30 runs the Python suite with DEVICE=cpu. After this PR and #30 merge, that pass runs these same cases on the CPU. The C++ test already runs the CPU and the GPU in one pass.

CONTRIBUTING checklist

• Formatting: uvx pre-commit run --files on the changed files (clang-format, black, isort).

• Tests: new C++ and Python tests, listed above.

• Benchmarks: no benchmark. The change adds a compare and a subtract to the signed-integer Divide only. Float and unsigned paths do not change. No local build was done (the laptop disk is shared). CI is the build and test authority.

• ☐ I understand it is strictly prohibited to use AI to write PR description

• AI usage disclosure: this change and this description were written by Claude (Opus 5.5) under David Tai's direction and review.

🤖 Generated with Claude Code