Sample randint directly instead of through a float32 uniform

[axiom-of-choice]

Proposed changes

Fixes #3926.

Problem

randint sampled through a float32 uniform and cast the result:

auto u = uniform(low, high, shape, float32, key, s);
return astype(maximum(u, low, s), dtype, s);

float32 has a 24-bit mantissa, so once the bounds reach 2^24 the integer interval is gone before the cast. Three distinct symptoms, all on main:

>>> key = mx.random.key(42)
>>> x = mx.random.randint(2**24, 2**24 + 2, (10_000,), dtype=mx.int32, key=key)
>>> sorted(set(x.tolist()))
[16777216, 16777218]      # returns `high` (excluded); never returns 2**24 + 1

>>> y = mx.random.randint(2**40, 2**40 + 1024, (10_000,), dtype=mx.int64, key=key)
>>> len(set(y.tolist()))
1                         # a 1024-wide interval collapses to one value

At 2^24 adjacent float32 values are two integers apart; at 2^40 the spacing is 2^17, so the whole interval sits between two representable floats.

Change

Sample the integers directly. The interval width is computed in int64 and reinterpreted as uint64, then 64 uniform bits (two bits() draws) are reduced onto it with a remainder.

The residual modulo bias is bounded by range / 2^64, below 1e-9 for any range under 2^32, and exactly zero when the width is a power of two. That is against the current behaviour of making values unreachable outright.

Empty intervals (high <= low) keep returning low, as before.

Relationship to #3936

#3936 proposes the same integer-domain approach, and I want to be clear that the idea is not mine alone; I reached it independently and said so there, offering the fix below as a patch to that branch rather than a competing PR. Opening this now so there is a reviewable version that closes the remaining case; happy for it to be closed in favour of #3936 if that branch picks the fix up.

The gap is that #3936 clamps the signed difference:

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

For an interval wider than 2^63 the true width overflows int64 and wraps negative, so maximum(range, 1) reads it as negative and clamps to 1. Every draw then lands on low.

I built that branch (6de5378, CPU-only, macOS/arm64) and measured 20,000 draws per interval against this one:

interval	#3936	this PR
[2^24, 2^24+2)	2 distinct ✅	2 distinct ✅
[2^40, 2^40+1024)	1024 ✅	1024 ✅
[0, 2^63-1)	20000 ✅	20000 ✅
[-2^62, 2^62)	1 ❌	20000 ✅
[-2^63, 2^63-1)	1 ❌	20000 ✅

# on #3936
>>> x = mx.random.randint(-(2**62), 2**62, (20_000,), 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 a constant, the same failure this work sets out to fix, relocated to a wider magnitude. An int64 interval can be up to 2^64 - 1 wide, and uint64 is the only type that holds every width randint can legitimately be asked for, which is why the clamp has to happen after the reinterpretation:

auto empty = less_equal(hi, lo, stream);      // signed, where the compare is meaningful
auto range = where(
    empty,
    array(1, uint64),
    astype(subtract(hi, lo, stream), uint64, stream),
    stream);

One smaller point: #3936's comment argues numpy remainder semantics keep the 64-bit path exact when the combined draw goes negative. The sign of the result is right, but remainder(raw, r) for negative raw is raw % r + r, which skews the distribution toward the low end of the interval by roughly 2^63 mod r. Reducing in uint64 removes the need for that argument.

Verification

Reported cases, and the overflow edges:

[2^24, 2^24+2)          in_bounds=True  distinct=2
[2^40, 2^40+1024)       in_bounds=True  distinct=1024
[-2^62, 2^62)           in_bounds=True  distinct=20000
[-2^63, 2^63-1)         in_bounds=True  distinct=20000
[0, 2^63-1)             in_bounds=True  distinct=20000

Uniformity over 600,000 draws. Every value reachable, worst-case deviation from uniform ~1%:

[0, 6)                  coverage=6/6      max_rel_dev=0.0036
[2^24, 2^24+2)          coverage=2/2      max_rel_dev=0.0003
[-10, 10)               coverage=20/20    max_rel_dev=0.0147

Existing behaviour preserved: randint(10, -10, ...) still returns low, same-key draws still reproduce, broadcasting of array bounds unchanged, and bool / uint8 dtypes unchanged.

Suites, CPU-only build on macOS/arm64:

• python/tests/test_random.py: 16 passed
• python/tests/: 727 passed, 66 skipped, 9718 subtests passed
• ./build-tests/tests/tests: 244 cases, 3245 assertions, all passed

I could not exercise the Metal path (no full Xcode locally, so no metal compiler). The change is in mlx/random.cpp above any backend dispatch and the issue reproduces identically on mx.cpu and mx.gpu, but the GPU path is the part worth a second look.

Tests added

• python/tests/test_random.py::test_randint_exact_integer_range: the 2^24 and 2^40 cases plus a width above 2^63.
• python/tests/test_random.py::test_randint_uniform_coverage: every value in a small interval is reachable at roughly equal frequency.
• tests/random_tests.cpp: the same three intervals in the existing test random randint case.

Both Python tests fail on main and pass with the change.

Checklist

• I have read the CONTRIBUTING document
• I have run pre-commit (on the changed files: clang-format, black, isort all pass) 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 (if needed). The documented contract ("integers in [low, high) with equal probability") is unchanged; this makes the implementation match it

[PhilipJohnBasile]

Replied in detail on #3936 (#3936 (comment)) rather than here, which left this PR with no response at all — correcting that.

Short version, so this thread stands on its own:

• Your report is confirmed on a real CPU-only build at 20,000 draws per interval. The wide-int64 collapse you found is real, your diagnosis of the signed clamp is right, and your fix is now on Fix mx.random.randint precision loss beyond 2^24 (float32 domain sampling) #3936's head.
• This PR has one remaining case, the same one Fix mx.random.randint precision loss beyond 2^24 (float32 domain sampling) #3936's head has, because both use less_equal(hi, lo) after the int64 cast: a uint64 interval that straddles 2^63 (low < 2^63 <= high) reads as empty and collapses to low. Its mirror image also bites — a genuinely inverted uint64 interval straddling 2^63 reads as non-empty and emits a spread instead of the contracted constant.
• I first proposed a one-line fix for it and have since retracted that: comparing the raw bounds promotes signed × uint64 to float32, which reintroduces the same precision bug in the predicate and additionally opens a modulo-by-zero on float bounds. Details and measurements are in the Fix mx.random.randint precision loss beyond 2^24 (float32 domain sampling) #3936 comment.
• What does survive both directions is comparing in the output domain:

auto cmp_dtype = issubdtype(dtype, unsignedinteger) ? uint64 : int64;
auto empty = less_equal(
    astype(high, cmp_dtype, stream), astype(low, cmp_dtype, stream), stream);

Happy either way on which PR lands it — yours with variant C folded in and I close #3936, or mine and you close this. Your reinterpret-as-uint64 core is the right approach regardless, and it's yours.