Keep randint samples inside [low, high) - #4012
Conversation
zcbenz
added
the
await verification
| // is not always representable in float32. Clamping to `low` last keeps an | ||
| // empty interval collapsing to `low`. | ||
| auto hi = astype(subtract(high, array(1, high.dtype()), s), dtype, s); | ||
| return maximum(minimum(out, hi, s), astype(low, dtype, s), s); |
There was a problem hiding this comment.
Should we pass high - 1 to uniform instead of using minimum(out, high - 1)?
There was a problem hiding this comment.
I tried that first and it loses the top value, because uniform(low, high - 1) samples [low, high - 1) and the floor then tops out at high - 2:
low=0 high=4 this PR -> [0, 1, 2, 3] high-1 arg -> [0, 1, 2]
low=-3 high=3 this PR -> [-3 .. 2] high-1 arg -> [-3 .. 1]
low=0 high=2 this PR -> [0, 1] high-1 arg -> [0]
The last row is the clearest, the two value case collapses to a constant.
There is a second reason not to push the bound through uniform. The argument would go in as float32, and high - 1 is not always representable there, so it can round back up to high or past it:
high-1 = 1073741823 -> float32 1073741824
high-1 = 2147483646 -> float32 2147483648
That is why the clamp is built in the integer domain instead.
One thing I should be straight about. I could not get floor(uniform(low, high)) to actually reach high in testing, six seeds at 200k samples across 2**24+1, 2**30, [-2**30, 2**30) and INT32_MAX, all stayed in range. So the minimum is defensive rather than something I can show firing. Happy to drop it if you would rather not carry a guard without a reproducer, though the floor change is the part that fixes the reported bug either way.
There was a problem hiding this comment.
Hmm I think we should just return floor(uniform(low, high)), since uniform is guarenteed to return between [low, high). I guess the bug was caused by astype rounding the result to high, and floor should have fixed it.
I wonder if we should remove maximum(u, low) too, I dont really see a way that we would get a number below low.
There was a problem hiding this comment.
I have to correct myself. Earlier I said I could not get floor(uniform(low, high)) to reach high. That was wrong, I was sweeping wide ranges, which is the wrong place to look. It reproduces immediately on a narrow range straddling the float32 integer limit:
randint(2**24, 2**24 + 2) without the upper clamp
int32 observed max 16777218 high = 16777218
uint32 observed max 16777218 high = 16777218
uniform is in [low, high) as a real interval, but high here is not representable in float32, so low + range * u rounds up to exactly high and the floor keeps it. With minimum(out, high - 1) the same 12 seeds stay in range. Those two cases are in this PR's test_randint, so dropping the clamp fails its own test.
On maximum(u, low), that one is load bearing too, and it is your existing tests that pin it rather than anything I added:
python/tests/test_random.py:233 mx.random.randint(10, -10, ...) asserts all == 10
tests/random_tests.cpp:509 "Check high < low => all equals to low"
I tried the bare version locally and both of those fail, randint(10, -10) returns values in [-10, 9] instead of all 10.
So I have left the code as it is. You are right that floor is what actually fixes the reported bug, the two clamps are only guarding the unrepresentable-bound case and the inverted-range contract. Happy to reshape or comment them differently if you would prefer that spelled out another way.
There was a problem hiding this comment.
Thanks for experimenting, I wonder if the algorithm proposed by JAX would be better:
https://docs.jax.dev/en/latest/_autosummary/jax.random.randint.html
def randint_via_uniform(key, shape, minval, maxval, dtype):
u = jax.random.uniform(key, shape, minval=minval - 0.5, maxval=maxval - 0.5)
return u.round().astype(dtype)Current approach feels like spending too much efforts patching around floating number details.
There was a problem hiding this comment.
I tried it, and it does not get us out of the floating point details, it lands in the same trap one step earlier.
On the exact case this PR is about, randint(2**24, 2**24 + 2), the rounding form is out of range on every seed I ran:
jax-style out-of-range seeds 12/12 values seen [16777216, 16777218]
this PR out-of-range seeds 0/12 values seen [16777216, 16777217]
16777218 is high, and 16777217 never appears at all, so it both escapes the interval and drops a valid value. The reason is that the shift is applied in float32 before sampling:
float32 spacing at 2**24 = 2.0
low - 0.5 = 16777215.5 -> float32 16777216.0
high - 0.5 = 16777217.5 -> float32 16777218.0
At that magnitude the half cannot be represented, so both bounds move up by 0.5 and round has nothing left to correct.
It also changes behavior your tests pin. randint(10, -10) and randint(20, -10) both stop returning low:
python/tests/test_random.py:233 randint(10, -10) all == 10 jax-style False
tests/random_tests.cpp:509 randint(20, -10) all == 20 jax-style False
On everything else the two agree exactly, same values and same counts over [-3, 3), so the difference really is only at the boundary.
I do take the point that this reads as patching around float behavior. That is fair, and if you want it out of randint entirely the honest fix is to draw integers rather than route through uniform, generating the range directly from bits and rejecting the overhang. Happy to do that instead if you would prefer it, though it is a bigger change than this bug needs.
There was a problem hiding this comment.
Thanks for trying, this is probably the best we can get with current algorithm and I'm good with it. I will do some research to have a fully understanding of the root cause and then merge.
zcbenz
removed
the
await verification
Rounding in the float32 uniform could land a sample on high itself, and the cast to the output dtype truncated towards zero instead of flooring, so low was unreachable for intervals spanning negative values and a bool output was almost always true. Take the floor before the cast and clamp to high - 1 in the integer domain, where the bound is exact.
d5fc467 to
f03da6c
Compare
2 participants
Fixes #3926.
mx.random.randintdraws a float32 uniform on[low, high)and casts it to the output dtype. Two things go wrong on the way.Rounding in float32 can land a sample on
highitself, so the excluded upper bound comes back:The cast also truncates towards zero instead of flooring, so for an interval that spans negative values
lowis never produced and zero collects twice the weight of every other value:That is the same reason
mx.random.randint(0, 1, dtype=mx.bool_)returnsTruealmost always: any nonzero float casts toTrue, and0is the only value the interval allows.The fix takes the floor before the cast and clamps the result to
high - 1. The clamp bound is computed in the integer domain rather than in float32, becausehigh - 1is not always representable there. At2**24the float32 spacing is already 2, so clamping in float would round the bound back ontohighand also drop2**24 + 1, which clamping after the cast keeps. The clamp tolowis applied after the clamp tohigh - 1so an empty interval still collapses tolow, which is the existing behaviour and is covered by a test.This keeps the uniform based algorithm and only clamps, per the direction on the issue, which is also what JAX documents for
jax.random.randint. Wide intervals and bounds beyond the float32 integer resolution still cannot reach every integer, for example[2**40, 2**40 + 1024)still collapses to a single value, so I added a note to the docstring saying so. Every sample is at least in range now.Verified on an M3 Pro running macOS 26, on both the CPU and the Metal stream.
Before, on 0.31.2:
After:
I extended
test_randintinpython/tests/test_random.pyand the C++test random randintcase with the bound, negative interval and boolean cases above.