Skip to content

Fix mx.random.randint precision loss beyond 2^24 (float32 domain sampling) - #3936

Open
PhilipJohnBasile wants to merge 6 commits into
ml-explore:mainfrom
PhilipJohnBasile:agent/fix-randint-float32-range
Open

Fix mx.random.randint precision loss beyond 2^24 (float32 domain sampling)#3936
PhilipJohnBasile wants to merge 6 commits into
ml-explore:mainfrom
PhilipJohnBasile:agent/fix-randint-float32-range

Conversation

@PhilipJohnBasile

@PhilipJohnBasile PhilipJohnBasile commented Jul 28, 2026

Copy link
Copy Markdown

Problem

mx.random.randint sampled 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 excluded high can 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:

  • an int64 interval can be up to 2^64 - 1 wide, so its width must be interpreted as uint64;
  • interval emptiness cannot be tested in one fixed domain: same-sign non-negative bounds compare in uint64, same-sign negative bounds compare in int64, and unlike signs determine the ordering without a cast;
  • the range divisor must be clamped independently of the predicate, because float bounds can saturate to the same integer and otherwise expose raw random bits through remainder(x, 0).

Change

  • Draw random bits directly and reduce them in the integer domain instead of routing through float32.
  • Combine two independent 32-bit draws for every output dtype. One 32-bit draw caused severe modulo bias for large uint32 intervals even though it could reach every value.
  • Interpret wrapped interval widths as uint64, including widths above int64_max.
  • Select the empty-interval comparison per element from the signs of both bounds: same-sign non-negative uses uint64, same-sign negative uses int64, and unlike signs are ordered directly.
  • Clamp the unsigned divisor to at least one independently of the predicate, preserving the existing empty/inverted-interval behavior of returning low.
  • Add Python coverage for float32 precision boundaries, widths above int64_max, uint64 intervals straddling 2^63, same-sign and cross-sign mixed int64/uint64 bounds, elementwise predicates, float saturation, negative bounds with unsigned output, determinism, small-range uniformity, and the large-uint32 modulo-bias case.
  • Add the core precision, wide-interval, and cross-sign cases to the C++ random test.

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 the 2^64 possible 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 uint64 framing, 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-commit on the three changed files: pass
  • python/tests/test_random.py: 20/20 passed
  • python -m unittest discover -v python/tests: 794 passed, 65 skipped
  • C++ test binary: 244/244 cases, 3248/3248 assertions passed

Metal was not exercised locally; backend CI and maintainer review remain required.

Checklist

  • Read CONTRIBUTING.md
  • Added regression tests that distinguish the fixed cases
  • Ran the repository format hooks on changed files
  • Ran the focused and full CPU test suites
  • Metal/backend CI

`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.
@PhilipJohnBasile
PhilipJohnBasile force-pushed the agent/fix-randint-float32-range branch from a390d4b to 6de5378 Compare July 29, 2026 00:08
@axiom-of-choice

Copy link
Copy Markdown

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 (6de5378, CPU-only, macOS/arm64) and sampled 20,000 draws per interval:

interval in bounds distinct values
[2^24, 2^24+2) yes 2 ✅
[2^40, 2^40+1024) yes 1024 ✅
[0, 2^63-1) yes 20000 ✅
[-2^62, 2^62) yes 1
[-2^63, 2^63-1) yes 1
>>> 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 == low

So randint(-2^62, 2^62) returns low for every sample — the same failure mode the PR fixes at 2^24, just relocated to a wider magnitude.

Cause. range is the signed difference:

auto range = subtract(hi, lo, stream);
auto safe_range = maximum(range, array(int64_t(1), int64), stream);

For [-2^62, 2^62) the true width is 2^63, which overflows int64 and wraps to -9223372036854775808. maximum(range, 1) then reads that as negative and clamps it to 1, so remainder(raw, 1) is always 0 and every draw lands on low.

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 randint can legitimately be asked for.

Two smaller notes, take or leave:

  • remainder on a negative dividend. The comment says numpy semantics make the 64-bit-combine path exact even when raw goes negative. That holds for the sign of the result, but remainder(raw, r) for negative raw is raw % r + r, which is not the same distribution over [0, r) as reducing the unsigned value — it skews toward the low end of the interval by roughly 2^63 mod r. Reducing in uint64 avoids needing the argument at all.

  • Test sensitivity. assertGreater(len(set(yv)), 500) on a 1024-wide interval passes even if a chunk of the interval is unreachable. Asserting == 1024 pins full coverage, which is the property that actually broke.

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 python/tests/test_random.py (16) and the full C++ suite (244 cases / 3245 assertions) with the above, if the uint64 framing is useful to compare against.

@axiom-of-choice

Copy link
Copy Markdown

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.
@PhilipJohnBasile

Copy link
Copy Markdown
Author

Thank you — this is a real bug and you were right to flag it rather than fork the effort. Reproduced, fixed in 3023040a, and I took all three of your points.

The width overflow. Confirmed exactly as you described. I checked the arithmetic independently before changing anything:

interval true width as int64 maximum(range, 1)
[-2^62, 2^62) 9223372036854775808 -9223372036854775808 1
[int64_min, int64_max) 18446744073709551615 -1 1

So every draw lands on low — the same failure this PR fixes at 2^24, just relocated, which is a fair thing to have caught in a patch whose whole point is that magnitude shouldn't matter.

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 lo + offset is done in uint64 rather than int64. For a full-width interval the offset can exceed int64_max, so the signed add would be overflow UB; unsigned wrap is well defined and the bit pattern reinterprets correctly into the target dtype.

Your remainder point. Agreed, and reducing unsigned removes the need for the argument entirely rather than making it more carefully. The old comment claimed numpy's negative-dividend semantics made the 64-bit path exact; as you say, raw % r + r is not uniform on [0, r) — it skews the low end by 2^63 mod r. An unsigned reduction is uniform by construction, so that comment is gone rather than reworded.

Your test point. Also right, and it's the sharper of the two secondary notes: assertGreater(len(set(yv)), 500) on a 1024-wide interval passes with half the interval unreachable, which is precisely the property that breaks. Now assertEqual(len(set(yv)), 1024).

Verified on macOS/arm64 (CPU), 20,000 draws per interval:

interval in bounds distinct
[2^24, 2^24+2) yes 2
[2^40, 2^40+1024) yes 1024
[0, 2^63-1) yes 20000
[-2^62, 2^62) yes 20000
[int64_min, int64_max) yes 20000

randint(-(2**62), 2**62, (20000,)) now gives min=-4610627987375572333 max=4611483732735255193, 20000 distinct, against a single repeated value before. 16/16 test_random tests pass.

New regression test is test_randint_interval_wider_than_int64_max, covering both of your intervals. Credited you in the commit message — say the word if you'd rather be listed differently, or if you'd prefer to carry the fix in your own PR, I'm happy to close this in favour of it.

@PhilipJohnBasile

PhilipJohnBasile commented Aug 1, 2026

Copy link
Copy Markdown
Author

Edit log. Corrected three times, always in place so no wrong version is left standing. (1) It first said I would fold @axiom-of-choice's fix in — it was already folded (3023040a); I'd read a stale gh pr diff instead of the head SHA. (2) It then proposed a one-line fix, which I had not built. I built it: it's wrong, and it's retracted below. (3) An audit of my own text found three further errors in it — the trigger condition, the reachability caveat, and the claim that this isn't a regression. All three are corrected below, and all three make the bug more serious than I first described.

@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 (MLX_BUILD_METAL=OFF) off 3023040a, 20,000 draws, key(7).

1. The trigger is straddling, not magnitude

I originally wrote "uint64 bounds at or above 2^63 collapse." That is the wrong generalization. The actual condition is low < 2^63 <= high — the interval has to straddle 2^63. Intervals with both bounds above 2^63 are computed correctly:

uint64 interval head 3023040a
[0, 2^64-1) straddles 1 distinct
[2^62, 2^64-1) straddles 1 distinct
[2^63, 2^64-1) both above 20000 ✅
[2^63+2^40, 2^63+2^41) both above 20000 ✅
[0, 2^62) both below 20000 ✅

The mechanism is unchanged: less_equal(hi, lo) runs after astype(high, int64) has wrapped, so a straddling interval is the one where exactly one endpoint flips sign and the ordering inverts.

2. There is a second bug, in the opposite direction

A genuinely empty (inverted) uint64 interval that straddles 2^63 reads as non-empty, so instead of the documented constant low it emits a full spread of out-of-contract values:

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 int64-only — there is no uint64 randint test at all, which is exactly why the first bug survived into the head too.

3. Correction: this is a regression on that interval

I wrote that main is "also wrong here, returning varied garbage," so the collapse wasn't a regression. That was wrong, and it understated the severity. For the interval I named, stock 0.32.0 is in-bounds and uniform:

stock 0.32.0, uint64 [0, 2^64-1), 20k draws:
  19992 distinct, in-bounds, deciles [2035, 1986, 1948, 1994, 2061, 1952, 2020, 2035, 1955, 2014]
head  3023040a, same interval:
      1 distinct, all zeros

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 [2^63+2^40, 2^63+2^41) main yields 2 distinct values and the head correctly yields 20000. Which is worse depends on the interval.)

Also, my reachability caveat was too narrow. Only Python int scalars above int64 max raise std::bad_cast; np.uint64 scalars pass straight through and collapse identically:

mx.random.randint(np.uint64(0), np.uint64(2**64 - 1), (20_000,), dtype=mx.uint64)  # -> all zeros

4. My proposed fix was wrong — do not apply it

I suggested less_equal(high, low) on the raw bounds and called it one line that changed nothing else. Built, it fixes the straddling collapse and breaks two other things, because MLX promotes signed × uint64 to float32:

mx.add(mx.array([0], mx.int64), mx.array([0], mx.uint64)).dtype  ->  float32
  • Mixed signed/unsigned bounds collapse. With a 24-bit mantissa, bounds closer than the float32 ULP round together and read as empty. At 2^62 the ULP is 2^39: low=int64(2^62), high=uint64(2^62+2^38) gives 20000 distinct on the head and 1 with my hoist.
  • Modulo by zero. randint validates only the output dtype, so float bounds are accepted. less_equal(0.9, 0.5) is false, so the interval reads non-empty while the truncated range is 0 — and MLX's remainder(x, 0) returns x unchanged rather than raising, so randint(0.5, 0.9, ...) emits raw PRNG bits: values spanning ±2^31 for an int32 output.

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 survives

The predicate has to be exact in both directions — comparing raw bounds is inexact for signed × uint64, comparing after the int64 cast is wrong for straddling intervals. Comparing in the output domain is exact for both:

auto cmp_dtype = issubdtype(dtype, unsignedinteger) ? uint64 : int64;
auto empty = less_equal(
    astype(high, cmp_dtype, stream), astype(low, cmp_dtype, stream), stream);
case A: head B: my hoist C: output domain
int32 [2^24, 2^24+2) 2 ✅ 2 ✅ 2 ✅
int64 [2^40, 2^40+1024) 1024 ✅ 1024 ✅ 1024 ✅
int64 [-2^62, 2^62) 20000 ✅ 20000 ✅ 20000 ✅
int64 [i64min, i64max) 20000 ✅ 20000 ✅ 20000 ✅
uint64 [0, 2^64-1) straddling 1 20000 ✅ 20000 ✅
uint64 [2^62, 2^64-1) straddling 1 20000 ✅ 20000 ✅
uint64 inverted straddling (empty) spread returns low
mixed i64/u64, gap 2^38 20000 ✅ 1 20000 ✅
float bounds [0.5, 0.9) returns 0 ✅ ±2^31 returns 0 ✅
[5,5) empty / [9,3) inverted low low low

All three arms pass the existing 16 tests in python/tests/test_random.py — the suite cannot distinguish them, which is the real finding here. Variant C should land with uint64 straddling, inverted-straddling, and mixed-dtype-bound regression tests.

6. Which PR

The straddling bug applies to #3955 as written too, since it uses the same less_equal(hi, lo) after the cast. I'd rather not push into your branch off the back of my own retracted suggestion, so your call: I put variant C plus tests here and you close #3955, or you take variant C into #3955 and I close #3936. No preference beyond it landing correctly — the reinterpret-as-uint64 core is yours either way.

@angeloskath — unrelated to #3937. State here: the wide-int64 bug is fixed on the head; two uint64 straddling bugs remain, with a verified fix not yet pushed pending the above.

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>
@PhilipJohnBasile

Copy link
Copy Markdown
Author

Pushed 0974b4010: variant C plus regression tests.

@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 uint64 bugs and a retracted fix recommendation of mine in the thread, and leaving that state overnight seemed worse than pushing something reviewable you can still override. If you'd rather this live in #3955, say so and I'll close #3936 in favour of it — the offer stands and I'll take whichever answer you give.

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 × uint64 to float32 and reintroduces this PR's own bug inside the predicate.

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 (AssertionError: 1 not greater than 19000), so it's a real regression test rather than a passing decoration. The mixed-dtype/float-bounds test passes on the previous head by design — it exists to stop anyone "fixing" this the way I first proposed.

Verification. CPU-only build (MLX_BUILD_METAL=OFF), 20,000 draws per interval:

  • test_random.py: 18/18 pass
  • test_ops.py + test_array.py: 243 pass, 19 skipped
  • behaviour matrix: all ten cases correct, including the two the previous head failed and the two my retracted one-liner would have broken

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 remainder/where on uint64 is unexercised by my testing.

…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>
@PhilipJohnBasile

Copy link
Copy Markdown
Author

Pushed f8997c8fc. The previous commit had a regression and I'd argued for it in this thread — correcting that.

An adversarial review of my own 0974b4010 found it: choosing the comparison domain from the output dtype means an unsigned output compares in uint64, so a negative low casts to a huge unsigned value, high <= low reads true, and the interval is treated as empty.

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 distinct

That 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 uint64 safe:

auto cmp_dtype = (issubdtype(low.dtype(), unsignedinteger) &&
                  issubdtype(high.dtype(), unsignedinteger))
    ? uint64
    : int64;

Otherwise int64 is exact for every signed bound and for every unsigned bound below 2^63 — which is exactly what the straddling case needs. So my earlier phrase "compare in the output domain" was the wrong description and the wrong code; the commit message says so.

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 0974b4010 and passes here.

Verification: CPU-only, 20,000 draws per interval. test_random 19, test_ops 146, test_array 97 — all pass. The straddling, inverted-straddling, mixed-dtype and float-bound cases all still behave as tabulated earlier in this thread; only the negative-low-unsigned row changed.

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:

  1. randint validates only the output dtype, so out-of-range or non-integer bounds ([127, 129) into int8, float bounds, bool) can still produce values outside [low, high).
  2. raw % range is not exactly uniform unless range divides 2^64 — the residual bias is range / 2^64, negligible for small ranges and not for very large ones. The "uniform by construction" phrasing in the surrounding comment overstates it.

Happy to address either here if reviewers want them in scope, or leave them for separate issues.

@PhilipJohnBasile

PhilipJohnBasile commented Aug 1, 2026

Copy link
Copy Markdown
Author

@axiom-of-choice — two things I owe you: a correction, and an answer to the question you asked twice.

The correction

In my reply above I wrote:

an unsigned reduction is uniform on [0, r) by construction

That is false, and I said it while contradicting you on a point you had already handled correctly. Reducing a uniform w-bit value modulo range is exactly uniform only when range divides 2^w.

For the 32-bit path with range = 3 * 2^30, residues in the low third have two preimages and the rest have one, so the low third receives probability 1/2 rather than 1/3, and the other two thirds get 1/4 each.

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 #3936

You 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:

  1. The empty-interval comparison must not use the already-int64-cast bounds. At 3023040a, [0, 2^64-1) collapses to all zeros, and an inverted interval straddling 2^63 returns a spread instead of the contracted low. The tested version selects the comparison domain from the output dtype — uint64 when it is unsigned, int64 otherwise. Struck: that describes 0974b4010, which is superseded. The current head f8997c8f derives the domain from the boundsuint64 only when both bounds are themselves unsigned — precisely because keying on the output dtype collapses randint(-2, 2, dtype=uint8), where the negative low casts to a huge unsigned value and the interval reads as empty. I inspected the older commit and published its behaviour as current; my apologies for the noise. Comparing the raw bounds without any cast is not safe either: MLX promotes a signed/uint64 pair to float32, so bounds closer together than the float32 ULP round into each other and read as empty.

  2. The final lo + offset should be unsigned. On a near-full-width interval the offset can exceed int64_max, so the signed add is overflow rather than the wrap it relies on; unsigned wrap is defined and preserves the intended bit pattern.

  3. The regression tests. Mine covered only signed output dtypes (int32 and int64), which is exactly why the uint64 regression survived them — the pre-existing 16-test suite passed on the broken head, on the broken follow-up, and on the corrected version, and could not distinguish the three. The straddling and inverted-straddling cases distinguish 3023040a from both later revisions; a negative-low-with-unsigned-output case is required to distinguish 0974b401 from f8997c8f.

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.

  1. The struck sentence above. f8997c8f, not 0974b4010, is the head; it keys on the bounds, not the output dtype.
  2. The straddling tests do not discriminate 0974b4010 from f8997c8f — both pass them. The case that separates those two revisions is a negative low with an unsigned output dtype. If piece 3 travels to Sample randint directly instead of through a float32 uniform #3955, that test should travel with it, otherwise it certifies less than it appears to.
  3. f8997c8f is still not exact. low = int64(2^62), high = uint64(2^64 - 1), dtype = uint64 selects int64, wraps high to -1, and reads the interval as empty. I first wrote this up as unbuilt analysis; it has since been measured on a CPU build of f8997c8f — all 20,000 draws return 2^62. Two other numbers from the same pass, since they bear on the pieces above: the 3 * 2^30 bias comes out 0.500469 / 0.249825 / 0.249706 against the predicted 1/2, 1/4, 1/4; and randint(-2, 2, dtype=uint8) returns a constant 254 on 0974b401 versus four values on f8997c8f, which is the case that separates those two revisions.

Also two things for whichever branch carries this: the source comment above the predicate in f8997c8f still says "Compare in the output domain", which is now the opposite of what the code does; and #3955 appears to have dropped #3936's broadcast-shape validation, so it can return a broadcast-expanded shape where historical randint raises. That is a fourth piece worth carrying over, and it is the kind of thing that is easier to notice now than after a merge.

(Second edit: two errors of my own in the first edit. "int64-only" was wrong — the tests covered int32 and int64, just no uint64 — and I had left a sentence claiming the straddling cases distinguish all three revisions standing above the later paragraph that contradicts it. Both corrected in place. The f8997c8f mixed-bounds case is now measured rather than argued.)

@axiom-of-choice

Copy link
Copy Markdown

@PhilipJohnBasile Thanks for the comments. I'm reading your comments and will write you a followup shortly.

@axiom-of-choice

Copy link
Copy Markdown

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.

@PhilipJohnBasile

Copy link
Copy Markdown
Author

Further review of f8997c8fc. Three things I wrote above are wrong, and the fix is less complete than I claimed — though, importantly, it is not a regression. Detail, because the reviewer-facing summary matters more than my bookkeeping.

What f8997c8fc actually does, measured. Fifteen intervals, CPU-only, 20,000 draws, key(7), against the branch base 3023040a:

failing regressions vs 3023040a fixed vs 3023040a
3023040a base 5
0974b4010 2 2 5
f8997c8fc (head) 2 0 3

So the head fixes three cases and introduces none; its two remaining failures were already failing on the base. 0974b4010 was the one that regressed, and that is what f8997c8fc corrected.

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 distinct

Corrections to what I wrote above:

  1. "Only when both bounds are themselves unsigned is uint64 safe"wrong, and it is why those two still fail. A non-negative int64 is exactly representable in uint64; the safe condition is that both bounds are non-negative in value, not that both dtypes are unsigned. My rule keys off dtype, so a non-negative signed low with a uint64 high above 2^63 falls to the int64 branch, where high wraps.

  2. "A negative bound with an unsigned output is out of contract, and the wrapping is unspecified"unsupported. The documented contract is equal-probability integers in [low, high) with no dtype-compatibility exception. I used that framing to justify a weak assertion in the test; the assertion is still the right one to make, but the justification was mine, not the docs'.

  3. "The residual bias is range / 2^64"wrong formula. The non-uniformity is real, the bound is not that.

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 float32 bounds that saturate to int64_max make range == 0 while the predicate reports non-empty, and remainder(x, 0) returns x, so the result is raw PRNG bits rather than a constant. The current head returns a constant there.

That points at the real defect underneath all of this, which none of my four attempts addressed: range and empty are computed independently, in different domains, and safe_range is only non-zero when they agree. Every version of this patch has been an attempt to make them agree by adjusting one side. The structural fix is to make disagreement harmless — clamp the unsigned range so it can never be zero — so a wrong predicate can at worst collapse an interval, never emit out-of-range bits.

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 uint64 reinterpretation is the part that has survived every round.

@axiom-of-choice

Copy link
Copy Markdown

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 int64 is representable in uint64, so low=int64(2^62) with high=uint64(2^64-1) falls to the int64 branch and high wraps. Working it out on paper first, neither domain is exact on its own and the split is clean:

interval true empty uint64 cmp int64 cmp
i64(2^62) × u64(2^64-1) no no ✅ yes
u64 [0, 2^64-1) no no ✅ yes
i64 [-2^62, 2^62) no yes no ✅
i64 [i64min, i64max) no yes no ✅
u64 [2^63+2^40, 2^62) yes yes ✅ no
[9, 3), [-2, -9) yes yes ✅ yes ✅

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 — astype(float32(2^63), int64) saturates to int64_max, so range goes to 0 while the predicate says non-empty, and remainder(x, 0) hands back x. But that's precisely what the clamp neutralises, and it's why I think your structural read is the right one rather than a fifth adjustment to one side. Make the divisor unconditionally ≥ 1 and a wrong predicate can only ever collapse an interval; it can't emit raw bits:

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, key(7) — your 15 intervals plus the float ones:

case f8997c8f + value predicate + clamp
i64(2^62) × u64(2^64-1) 1 20000 ✅
i64(0) × u64(2^64-1) 1 20000 ✅
u64 straddling ×2 20000 ✅ 20000 ✅
u64 inverted straddling low low
i64 [-2^62, 2^62), [i64min, i64max) 20000 ✅ 20000 ✅
i32 [2^24, 2^24+2) / i64 [2^40, +1024) 2 / 1024 ✅ 2 / 1024 ✅
randint(-2, 2, uint8) 4 ✅ 4 ✅
mixed gap 2^38 20000 ✅ 20000 ✅
[5,5) / [9,3) low low
f32 [1e19, 2e19) (saturates) low low
[0.5, 0.9) 0 ✅ 0 ✅

15/15. u64 [0, 2^64-1) over 200k draws: deciles within 0.9% of flat, all in bounds. python/tests/ 728 passed / 66 skipped / 9718 subtests, test_random.py 20/20, C++ 244 cases / 3245 assertions.

Since sign is per-element I checked the elementwise paths too, which is where I'd expect a where on a broadcast predicate to break: array bounds mixing signs per element, and one empty column alongside a non-empty one, both come out right. Small dtypes and bool bounds unchanged. One thing I hit that is not ours: randint(0, 2**64-1, dtype=uint64) with Python int bounds raises std::bad_cast on your head too, so it predates this work — passing mx.array(..., mx.uint64) is the way in.

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 r is bounded by r/2^64 only when r divides 2^64; in general it's the ratio between the largest and smallest bucket, ceil(2^64/r) / floor(2^64/r) - 1, which for any r below 2^32 is still under 1e-9. My PR description has your original formula in it and inherits the error.

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.

@PhilipJohnBasile

Copy link
Copy Markdown
Author

Pushed 6c0731933 with the combined follow-up we converged on.

@axiom-of-choice — I took your value-based split literally and kept it elementwise: both bounds non-negative compares in uint64; either negative compares in int64. The unsigned range is then clamped to at least 1 independently of the predicate, so saturated float bounds cannot expose raw bits through remainder(x, 0).

Coverage added in the same commit:

  • the two remaining mixed int64/uint64 intervals (low = 0 and 2^62, high = 2^64 - 1);
  • one broadcast row containing a non-empty wide interval beside an inverted interval, exercising the predicate per element;
  • float [0.5, 0.9) and saturating float32 [1e19, 2e19) bounds;
  • coarse uniformity/reachability over three small intervals;
  • your three C++ cases at 2^24, 2^40, and width 2^63.

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:

  • format hooks: pass
  • test_random.py: 20/20
  • full Python discovery: 794 passed, 65 skipped
  • C++: 244/244 cases, 3245/3245 assertions
  • warnings-as-errors build: pass

I did not exercise Metal locally. The branch is ready for your review and backend CI.

@PhilipJohnBasile

Copy link
Copy Markdown
Author

Follow-up pushed as 243dc2fc4 after an independent review caught two gaps in the first structural pass:

  1. “Either bound negative → compare as int64” is not exact for a cross-sign pair such as low=int64(-1), high=uint64(2^63). The predicate now classifies both signs: same-sign negative compares in int64, same-sign non-negative compares in uint64, and unlike signs determine the ordering without a cast. I added Python and C++ coverage for both cross-sign orderings.
  2. Non-64-bit outputs still used one 32-bit raw draw. For uint32 [0, 3*2^30), modulo reduction consequently put about half the samples in the first third instead of one third. Every output dtype now combines two independent 32-bit draws; a deterministic 200k-sample three-band test pins that regression.

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:

  • format hooks and git diff --check: pass
  • test_random.py: 20/20
  • full Python discovery: 794 passed, 65 skipped
  • warnings-as-errors CPU build: pass
  • C++: 244/244 cases, 3248/3248 assertions

Metal remains untested locally; backend CI and maintainer review are still required.

@axiom-of-choice

Copy link
Copy Markdown

Reviewed 243dc2fc4 on Metal. I installed Xcode + the Metal toolchain to close the GPU gap we'd both been declaring, so this is the first run of any of this on real hardware.

GPU results — 16/16, and bit-identical to CPU. M2 Pro, macOS 26.5.2, 20,000 draws, key(7), every interval run on both mx.gpu and mx.cpu and the tuples (distinct, min, max) compared:

interval gpu distinct gpu == cpu
i64(-1) × u64(2^63) cross-sign 20000 yes
u64(2^63) × i64(-1) inverted low yes
i64(2^62) × u64(2^64-1) 20000 yes
i64(0) × u64(2^64-1) 20000 yes
u64 [0, 2^64-1) straddling 20000 yes
u64 inverted straddling low yes
i64 [-2^62, 2^62) / [i64min, i64max) 20000 yes
i32 [2^24, 2^24+2) 2 yes
i64 [2^40, +1024) 1024 yes
randint(-2, 2, uint8) 4 yes
mixed gap 2^38 20000 yes
f32 [1e19, 2e19) saturating low yes
[0.5, 0.9) / [5,5) / [9,3) low yes

uint32 [0, 3*2^30) bands: [66770, 66866, 66364] on GPU, identical on CPU — your entropy fix holds on both backends.

Suites with Metal enabled: test_random.py 20/20, full Python discovery 749 passed / 45 skipped / 10584 subtests (vs 728/66 CPU-only, so 21 Metal-only tests also pass), C++ 260 cases / 3500 assertions (vs 244/3248). Warnings-as-errors build clean.

I also tried to break the parts that had never been exercised, and couldn't:

  • Independence of the two draws — over 200k samples of u64 [0, 2^64-1): corr(hi_half, lo_half) = -0.0015, hi == lo in 0 samples, both half-means within 0.1% of 0.5. Your key split is doing its job.
  • Per-bit uniformity — all 64 bit positions within 1% of 0.5.
  • Global seed() reproducibility and state advance; shape=(), (0,), (2,0,3); bool output ([0,2) gives 49958/50042, [0,1) all False); int8/uint8 with a range wider than the dtype domain.

One thing worth flagging: cost

The unconditional second draw is much more expensive on GPU than on CPU, because it pulls the whole pipeline into 64-bit:

randint uniform (old impl) ratio
1M int32, GPU 1.76 ms 0.29 ms 6.2x
10M int32, GPU 16.17 ms 2.73 ms 5.9x
10M int32, CPU 99.67 ms 43.74 ms 2.3x

Isolating the ops at 10M on GPU: astype u32→u64 0.94 ms ×2, shift+or combine 3.30 ms, remainder in uint64 4.71 ms, add in uint64 1.13 ms, final cast 0.94 ms — sums to ~14.8 ms of the measured 16.3 ms. The interesting part is that uint64 remainder costs 6.9x its uint32 counterpart (4.71 vs 0.69 ms) while uint64 divide and multiply are only ~1.2–1.8x. So it's 64-bit modulo specifically on Metal, not 64-bit width in general.

But I have to walk back the fix I was going to suggest. I was going to propose gating the second draw on range > 2^28 and keeping a 32-bit path below it. I worked out the actual bias before writing it and that threshold is wrong — the imbalance goes as r/2^32, and powers of two are exact while their neighbours are not:

range 32-bit bias 64-bit bias
10^4 2.3e-06 4.4e-16
10^6 2.3e-04 5.4e-14
2^20 0 0
2^20 + 1 2.4e-04 5.7e-14
2^24 0 0
2^24 + 1 3.9e-03 9.1e-13

At 2^24 + 1 a single draw is already off by 0.4%, so a 2^28 gate would have shipped a visible bias on very ordinary ranges. Only ranges below ~10^4, plus exact powers of two, are safe with one draw — which is a much narrower window than I claimed, and means your unconditional second draw is justified far lower than I'd assumed. Any fast path would have to key off range being a power of two or genuinely small, which is more complexity than the win probably warrants.

So I'd land this as it stands. The cost is a real regression for large int32 draws on GPU and worth a follow-up issue, but it buys correctness that the old path didn't have, and the alternative I had in mind was wrong. Flagging the numbers so it's a maintainer's decision rather than a surprise. #3955 is superseded by this — I'll close it in favour of this branch.

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 243dc2fc4 were yours alone. Your branch is the right one to land.

@axiom-of-choice

axiom-of-choice commented Aug 3, 2026

Copy link
Copy Markdown

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 mx.gpu vs mx.cpu with the tuples compared, plus test_random.py, full Python discovery (749/45 with Metal vs 728/66 CPU-only), and the C++ suite (260 cases / 3500 assertions vs 244/3248). Turnaround is a few minutes once the build is warm.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants