Fix mx.random.randint precision loss beyond 2^24 (float32 domain sampling) - #3936
Fix mx.random.randint precision loss beyond 2^24 (float32 domain sampling)#3936PhilipJohnBasile wants to merge 6 commits into
Conversation
`randint` sampled by drawing a float32 uniform value in [low, high) and rounding it to the target integer dtype. float32 has a 24-bit mantissa, so once |low| or |high| reaches 2**24 the rounding is no longer exact: a draw can round to the excluded `high`, skip valid integers entirely, or -- when the whole interval sits between two consecutive representable floats -- collapse every draw in the interval to a single constant, regardless of how narrow the requested interval is (e.g. randint(2**40, 2**40+1024, ...) returning one value for every draw). Sample directly in integer space instead: draw raw random bits and reduce them modulo the (int64) range. For dtypes whose own domain exceeds 2**32 (int64/uint64) combine two independent 32-bit draws into 64 bits of entropy; every other dtype's domain fits within one uint32 draw. `remainder` is numpy-style (result in [0, b), sign of the divisor), so this stays exact even when the combined 64-bit value is negative in two's complement. Existing behavior is preserved: broadcasting low/high, the `high <= low` degenerate case returning `low`, and per-dtype output shape/dtype. Verified against the exact reproduction from the reported issue on both CPU and Metal, plus the full existing test_random.py suite (15/15) and the broader suite of every other file that exercises randint (254 passed, 9460 subtests). Fixes ml-explore#3926.
a390d4b to
6de5378
Compare
|
I hit the same issue independently and arrived at the same approach — integer-domain sampling with a modulo reduction is, I think, clearly the right shape here. While testing my version I found one case this patch still gets wrong, and it looks like a one-line change, so flagging it rather than opening a competing PR. Interval widths above 2^63 collapse to a single value. Built this branch (
>>> x = mx.random.randint(-(2**62), 2**62, (20000,), dtype=mx.int64)
>>> min(x.tolist()), max(x.tolist()), len(set(x.tolist()))
(-4611686018427387904, -4611686018427387904, 1) # every draw == lowSo Cause. auto range = subtract(hi, lo, stream);
auto safe_range = maximum(range, array(int64_t(1), int64), stream);For The subtraction itself is fine — the wrapped bit pattern is the correct unsigned width. What breaks is comparing it as signed. Reinterpreting as uint64 before the clamp fixes it, with the empty-interval check done on the signed values where the comparison is still meaningful: auto empty = less_equal(hi, lo, stream);
auto range = where(
empty,
array(1, uint64),
astype(subtract(hi, lo, stream), uint64, stream),
stream);
auto offset = astype(remainder(u, range, stream), int64, stream);An int64 interval can be up to 2^64 - 1 wide, so uint64 is the only type that holds every width Two smaller notes, take or leave:
Happy to send these as a patch to your branch if that's easier than folding them in — whatever gets it merged. My local version passes |
|
Opened #3955 with the uint64 framing so there's a reviewable version that closes the >2^63 case, rather than leaving the finding sitting in a comment. It credits this PR for the approach and links back here. To be explicit about intent: I'm not trying to race you. If you'd rather fold the change into this branch, I'll close #3955 — the one-line version is in the comment above and in that PR's description. Whichever lands is fine by me. |
…lapse The width of an int64 interval does not itself fit in int64. `[-2**62, 2**62)` is 2**63 wide and `[int64_min, int64_max)` is 2**64 - 1 wide, so computing the range as a signed subtraction wraps negative, and the signed `maximum(range, 1)` clamp then reads that as negative and yields 1 -- so `remainder(raw, 1)` is always 0 and every draw collapses onto `low`. That is the same defect this PR fixes at 2**24, relocated to a wider magnitude. Reported by @axiom-of-choice, who hit it independently and supplied the fix rather than opening a competing PR. The wrapped bit pattern IS the correct unsigned width, so reinterpret the range as uint64 and do the empty-interval test on the signed values, where the comparison is still meaningful. Drawing and reducing the raw bits in uint64 also removes the need for the negative-dividend `remainder` argument: numpy semantics give `raw % r + r` for negative `raw`, which skews the low end of the interval by `2**63 mod r`, whereas an unsigned reduction is uniform on [0, r) by construction. The final add is done in uint64 as well, so the wrap for full-width intervals is well defined rather than signed overflow UB; the resulting bit pattern reinterprets correctly into the target dtype. Tests: new test_randint_interval_wider_than_int64_max covers both reported intervals. Also tightened the existing 1024-wide check from `assertGreater(len(set(yv)), 500)` to `assertEqual(..., 1024)` -- the loose form passes even when half the interval is unreachable, which is exactly the property that breaks. Verified on macOS/arm64 CPU: all five intervals in the report now in-bounds with full spread (20000/20000 distinct on the wide ones), 16/16 test_random tests pass.
|
Thank you — this is a real bug and you were right to flag it rather than fork the effort. Reproduced, fixed in The width overflow. Confirmed exactly as you described. I checked the arithmetic independently before changing anything:
So every draw lands on I used your uint64 reinterpretation with the empty-interval test on the signed values, and extended it one step further: the raw bits are now drawn and combined in uint64 too, and the final Your Your test point. Also right, and it's the sharper of the two secondary notes: Verified on macOS/arm64 (CPU), 20,000 draws per interval:
New regression test is |
@axiom-of-choice — summary: your report is confirmed on a real build; my proposed one-liner was wrong, don't apply it; the head has a second bug in the opposite direction that neither of us mentioned; and there's a fix that handles all of it. All numbers below: CPU-only builds ( 1. The trigger is straddling, not magnitudeI originally wrote "
The mechanism is unchanged: 2. There is a second bug, in the opposite directionA genuinely empty (inverted) lo, hi = mx.array([2**63 + 2**40], mx.uint64), mx.array([2**62], mx.uint64) # high < low: empty
mx.random.randint(lo, hi, (20_000,), dtype=mx.uint64, key=mx.random.key(7))
# head 3023040a: 20000 distinct values starting at 337491463012850 (contract says: all == low)Same wrapped predicate, opposite sign. Neither of us caught it because the tests added in this PR are 3. Correction: this is a regression on that intervalI wrote that It's coarse on main — the float32 path can only reach ~2^24 of the 2^64 values — but at this sample size it is in-bounds and flat, not garbage. So the head is a genuine regression there. (The reverse also holds elsewhere: for Also, my reachability caveat was too narrow. Only Python int scalars above int64 max raise mx.random.randint(np.uint64(0), np.uint64(2**64 - 1), (20_000,), dtype=mx.uint64) # -> all zeros4. My proposed fix was wrong — do not apply itI suggested
That first case works correctly today precisely because the current code compares after casting to int64, where it's exact. My hoist would have reintroduced, inside the emptiness test, the very class of bug this PR exists to remove. 5. A fix that survivesThe predicate has to be exact in both directions — comparing raw bounds is inexact for signed × auto cmp_dtype = issubdtype(dtype, unsignedinteger) ? uint64 : int64;
auto empty = less_equal(
astype(high, cmp_dtype, stream), astype(low, cmp_dtype, stream), stream);
All three arms pass the existing 16 tests in 6. Which PRThe straddling bug applies to #3955 as written too, since it uses the same @angeloskath — unrelated to #3937. State here: the wide- |
The empty-interval predicate was evaluated on the int64 casts of the bounds, which is not an exact domain for uint64. An interval with `low < 2**63 <= high` has its `high` wrap negative, so the ordering inverts, a valid interval reads as empty, safe_range clamps to 1, and every draw collapses onto `low`. The mirror image bites as well: a genuinely inverted interval straddling 2**63 reads as non-empty and returns a spread instead of the contracted constant `low`. Comparing the original bounds instead is also wrong. MLX promotes a signed/uint64 pair to float32, so the predicate would run with a 24-bit mantissa and round widely-separated bounds onto each other, reproducing this patch's own defect inside the emptiness test. It additionally opens a modulo by zero: randint validates only the output dtype, so float bounds reach the sampler, and remainder by 0 returns the dividend rather than raising. Compare in the output domain, which is exact for the dtype requested. Measured CPU-only at 20,000 draws per interval: uint64 [0, 2**64-1) and [2**62, 2**64-1) go from 1 distinct value to 20000; inverted straddling intervals return `low` again; and mixed int64/uint64 bounds 2**38 apart at 2**62, which the previous predicate handled correctly, still yield 20000 rather than collapsing. Intervals with both bounds above 2**63 were never affected and are covered as controls. Found while reviewing @axiom-of-choice's fix for the int64 width overflow, whose reinterpret-as-uint64 approach this builds on. The existing suite did not distinguish any of these variants, so the new tests cover the straddling, inverted-straddling, mixed-dtype and float-bound cases. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Pushed @axiom-of-choice — I said I'd let you pick which PR carries this and then pushed anyway, so to be straight about why: the head was sitting with two known The change. The empty-interval predicate now compares in the output dtype's domain: auto cmp_dtype = issubdtype(dtype, unsignedinteger) ? uint64 : int64;
auto empty = less_equal(
astype(high, cmp_dtype, stream), astype(low, cmp_dtype, stream), stream);Neither obvious alternative is exact, which is the whole point of the commit message: testing the int64 casts is wrong for intervals straddling 2^63 (in both directions), and testing the raw bounds promotes signed × Tests. Two added, covering straddling, inverted-straddling, mixed-dtype bounds and float bounds, with both-above-2^63 and both-below intervals as controls. The straddling test fails on the previous head ( Verification. CPU-only build (
The existing 16-test suite passed on the previous head, on my broken hoist, and on this — it could not tell the three apart. That's why the new tests are scoped the way they are. Not claimed: CPU only. I have not run this on Metal, and the GPU path for |
…tput The previous commit chose the comparison domain from the output dtype, so an unsigned output compared in uint64. A negative `low` then casts to a huge unsigned value, `high <= low` reads true, and the interval is treated as empty -- `randint(-2, 2, dtype=uint8)` collapsed from four distinct values to the constant 254. That is the same "looks like data" failure the rest of this patch exists to remove, and it was a regression: the prior behaviour returned a varied (if out-of-contract) sample. Only when both bounds are themselves unsigned is uint64 a safe domain, since only then can no operand be negative. Otherwise int64 is exact for every signed bound and for every unsigned bound below 2^63, which is what the uint64-straddling case needs. A negative bound with an unsigned output dtype is out of contract regardless, and its wrapping behaviour is unspecified, so the added test asserts only the property that regressed -- that the sample does not collapse to a constant. Found by an adversarial review of the previous commit. Verified CPU-only: test_random 19, test_ops 146, test_array 97 pass; the new test fails on the previous commit and passes here; the uint64 straddling, inverted-straddling, and mixed-dtype cases all still behave. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
Pushed An adversarial review of my own mx.random.randint(-2, 2, (20_000,), dtype=mx.uint8, key=mx.random.key(7))
# 3023040a (before my fix): [0, 1, 254, 255] 4 distinct
# 0974b4010 (my fix): [254] 1 distinct <- regression
# f8997c8fc (this): [0, 1, 254, 255] 4 distinctThat is precisely the failure mode this PR exists to remove — a silent collapse to a plausible-looking constant — and I reintroduced it while fixing the same class of bug one layer up. It is also the second time in this thread I've done that: variant B did it via float32 promotion, and I retracted it. Same lesson, one level deeper. The fix: the domain has to come from the bounds, not the output. Only when both bounds are themselves unsigned can no operand be negative, so only then is auto cmp_dtype = (issubdtype(low.dtype(), unsignedinteger) &&
issubdtype(high.dtype(), unsignedinteger))
? uint64
: int64;Otherwise A negative bound with an unsigned output is out of contract however you slice it, and the wrap behaviour is unspecified, so the added test asserts only the property that regressed: it must not collapse to a constant. It fails on Verification: CPU-only, 20,000 draws per interval. Two pre-existing gaps this PR does not close, flagged so nobody reads the tables above as a completeness claim — both predate it and are unchanged by it:
Happy to address either here if reviewers want them in scope, or leave them for separate issues. |
|
@axiom-of-choice — two things I owe you: a correction, and an answer to the question you asked twice. The correctionIn my reply above I wrote:
That is false, and I said it while contradicting you on a point you had already handled correctly. Reducing a uniform For the 32-bit path with Switching from a signed value to its unsigned bit pattern removes the negative-dividend mapping I was actually answering, but it does not remove modulo bias. Your #3955 documents that residual bias explicitly. I treated it as something you had missed when you had already priced it in. I had that backwards. Consolidate in #3955; I'll close #3936You offered twice to close yours. Looking at where the two branches actually ended up, yours is the better base:
Three pieces from #3936 are worth carrying over:
One caveat I'd rather state than have you inherit: selecting the domain this way fixes the cases we have tested, but it is not a proof that every mixed signed/unsigned or out-of-output-range bound combination has a well-defined contract. Bounds validation is a separate gap in both PRs. I can open those three against your branch, or hand you the diffs if you'd rather apply them yourself. You found the wide-interval overflow independently, brought a fix instead of a competing PR, and documented the bias I then told you about. Unless you'd rather it went the other way, I'll move those pieces to #3955 and close this. Edit — three further corrections, after a re-check against the actual head rather than the commit I had open.
Also two things for whichever branch carries this: the source comment above the predicate in (Second edit: two errors of my own in the first edit. " |
|
@PhilipJohnBasile Thanks for the comments. I'm reading your comments and will write you a followup shortly. |
|
Works for me, let's land it on yours. I reproduced both uint64 cases on my branch, so #3955 has them too. And variant C is the right call: comparing the raw bounds was never going to hold once signed × uint64 promotes to float32. Your branch is ahead on all of this, so I'll close #3955 rather than keep two half-fixes in flight. Two things from mine worth carrying over if you want them:
I'll send both as a patch to your branch. Shout if you'd rather pull the diff off #3955 yourself. |
|
Further review of What
So the head fixes three cases and introduces none; its two remaining failures were already failing on the base. The two cases still failing, both mixed-dtype bounds: lo, hi = mx.array([2**62], mx.int64), mx.array([2**64 - 1], mx.uint64) # also (0, 2**64-1)
mx.random.randint(lo, hi, (20_000,), dtype=mx.uint64, key=mx.random.key(7)) # -> 1 distinctCorrections to what I wrote above:
Why I am not pushing the obvious fix. Switching the predicate to a value-based non-negativity test does close both remaining cases — 15/15, full suites green. It also opens a worse hole: non-negative That points at the real defect underneath all of this, which none of my four attempts addressed: I would rather leave the head where it is (strictly better than the base, no regressions, two known gaps documented here) than land a fifth iteration tonight. Each of the previous four passed my own case set and was broken by review afterwards; that is a pattern about my case sets, not about the individual bugs. Happy to do the clamp plus the value-based predicate as a follow-up with the mixed-bound and float-bound intervals as tests — or to hand the whole predicate to @axiom-of-choice, whose |
|
I'll take the predicate, and I think the clamp is the answer — I built both together and the float hole you were worried about doesn't open. Your correction (1) is the key one: keying off dtype instead of value is exactly why the two mixed-bound cases survive. A non-negative
Both bounds non-negative → uint64 is exact; either negative → int64 is exact. Nothing else is. On the float saturation: you're right that a value-based predicate alone reintroduces it — auto safe_range = maximum(
where(empty, array(uint64_t(1), uint64), astype(range, uint64, stream), stream),
array(uint64_t(1), uint64), stream);Measured, CPU-only, 20,000 draws,
15/15. Since sign is per-element I checked the elementwise paths too, which is where I'd expect a Two of your three corrections I'd only partly agree with, for what it's worth: (2) I'd keep the weak assertion, but agree the docs don't license it, so the comment should say "unspecified here" rather than "out of contract". (3) agreed the formula is wrong — the bias for a range So: happy to push the predicate + clamp + the mixed-bound and float-bound tests to your branch as a second commit, on top of the C++/uniformity one. Your call on ordering — I can hold it until you've looked at the first, or send them together. Also worth saying: four rounds of this and the thing that keeps finding bugs is that we're both building and measuring rather than reasoning about it. Your case sets weren't the problem so much as that no single case set was going to be enough here. |
|
Pushed @axiom-of-choice — I took your value-based split literally and kept it elementwise: both bounds non-negative compares in Coverage added in the same commit:
I also corrected the code comment and PR body: unsigned modulo avoids the signed-remainder skew, but it does not make a non-power-of-two range mathematically exact. The description now states the one-raw-draw bucket imbalance rather than the earlier incorrect bias formula. Fresh CPU-only verification on macOS/arm64:
I did not exercise Metal locally. The branch is ready for your review and backend CI. |
|
Follow-up pushed as
I also corrected the description's modulo-bias limit: 64 raw bits make the bias negligible for <=32-bit ranges, but a near-full-width 64-bit range can still have a 2x relative bucket difference. Exact uniformity would require rejection sampling and is not claimed here. Fresh CPU-only verification after the follow-up:
Metal remains untested locally; backend CI and maintainer review are still required. |
|
Reviewed GPU results — 16/16, and bit-identical to CPU. M2 Pro, macOS 26.5.2, 20,000 draws,
Suites with Metal enabled: I also tried to break the parts that had never been exercised, and couldn't:
One thing worth flagging: costThe unconditional second draw is much more expensive on GPU than on CPU, because it pulls the whole pipeline into 64-bit:
Isolating the ops at 10M on GPU: But I have to walk back the fix I was going to suggest. I was going to propose gating the second draw on
At So I'd land this as it stands. The cost is a real regression for large For the record on attribution, since you've been careful about it: the uint64 reinterpretation and the width-overflow report were mine, the value-based split and the clamp came out of your review, and the cross-sign and entropy gaps in |
|
Closed #3955 in favour of this branch — #3955 (comment) records why, with the GPU numbers repeated there so the two threads don't disagree. One standing offer, since it's the caveat that's been on every comment in this thread from both of us: I have Xcode 26.6 and the Metal toolchain set up now, so the "Metal not exercised locally" line no longer applies on my side. If you push further revisions here, or if a maintainer asks for backend coverage, ping me and I'll re-run the full GPU matrix — 16 intervals on |
Problem
mx.random.randintsampled through a float32 uniform and cast the result to the requested integer dtype. Once a bound reaches the float32 precision boundary, valid integers become unreachable, the excludedhighcan be returned, or a narrow interval can collapse to one value.Moving the sampler into integer space exposed three related edge cases that this branch now handles together:
int64interval can be up to2^64 - 1wide, so its width must be interpreted asuint64;uint64, same-sign negative bounds compare inint64, and unlike signs determine the ordering without a cast;remainder(x, 0).Change
uint32intervals even though it could reach every value.uint64, including widths aboveint64_max.uint64, same-sign negative usesint64, and unlike signs are ordered directly.low.int64_max, uint64 intervals straddling2^63, same-sign and cross-sign mixedint64/uint64bounds, elementwise predicates, float saturation, negative bounds with unsigned output, determinism, small-range uniformity, and the large-uint32modulo-bias case.Unsigned remainder avoids the asymmetric negative-dividend path. It still has the ordinary modulo bias when the interval width does not divide
2^64; each bucket differs by at most one of the2^64possible raw draws. That difference is negligible for ranges up to 32 bits, but it can still be a 2x relative difference for a near-full-width 64-bit interval. This PR does not claim rejection-sampling exactness.Collaboration
@axiom-of-choice independently found the wide-range failure, contributed the
uint64framing, proposed the C++ and uniformity coverage, and validated the final value-based predicate plus structural clamp. #3955 contains the comparison implementation and investigation history; we agreed to carry the combined fix here.Verification
CPU-only macOS/arm64 build (
MLX_BUILD_METAL=OFF, warnings treated as errors):pre-commiton the three changed files: passpython/tests/test_random.py: 20/20 passedpython -m unittest discover -v python/tests: 794 passed, 65 skippedMetal was not exercised locally; backend CI and maintainer review remain required.
Checklist
CONTRIBUTING.md