diff --git a/mlx/random.cpp b/mlx/random.cpp index def3169cb5..8f27425299 100644 --- a/mlx/random.cpp +++ b/mlx/random.cpp @@ -279,8 +279,49 @@ array randint( throw std::invalid_argument( "[randint] randint only accepts integer dtypes and bool."); } - auto u = uniform(low, high, shape, float32, key, s); - return astype(maximum(u, low, s), dtype, s); + auto stream = to_stream(s); + + // Sample the integers directly rather than by casting a float32 uniform. + // Routing through float32 loses the integer interval whenever the bounds are + // not exactly representable: at 2^24 consecutive float32 values are already + // two apart, so values become unreachable, `high` itself can be returned, + // and a wide int64 interval can collapse to a single value. + auto lo = astype(low, int64, stream); + auto hi = astype(high, int64, stream); + + // Width of the half-open interval, as uint64: an int64 interval can be up to + // 2^64 - 1 wide, which only uint64 holds. The subtraction is done in int64 + // and reinterpreted, which is correct even when the difference overflows + // int64 — the wrapped bit pattern is the intended unsigned width. Note the + // clamp below therefore has to act on the *unsigned* value: clamping the + // signed difference would read a genuine width above 2^63 as negative and + // collapse the whole interval onto `low`. + // + // An empty interval (high <= low) is documented to return `low`, and a + // remainder by zero is undefined, so those entries get a divisor of 1, which + // yields an offset of 0. + auto empty = less_equal(hi, lo, stream); + auto range = where( + empty, + array(1, uint64), + astype(subtract(hi, lo, stream), uint64, stream), + stream); + + // 64 uniform bits per sample, assembled from two 32-bit draws. + auto wide_shape = shape; + wide_shape.push_back(2); + auto words = astype(bits(wide_shape, 4, key, stream), uint64, stream); + auto halves = split(words, 2, -1, stream); + auto u = bitwise_or( + left_shift(squeeze(halves[0], -1, stream), array(32, uint64), stream), + squeeze(halves[1], -1, stream), + stream); + + // Reduce onto [0, range). The residual modulo bias is bounded by + // range / 2^64: below 1e-9 for any range under 2^32, and exactly zero + // whenever range is a power of two. + auto offset = astype(remainder(u, range, stream), int64, stream); + return astype(add(lo, offset, stream), dtype, stream); } array bernoulli( diff --git a/python/tests/test_random.py b/python/tests/test_random.py index 551c32993c..7f1f4a8633 100644 --- a/python/tests/test_random.py +++ b/python/tests/test_random.py @@ -234,6 +234,41 @@ def test_randint(self): mx.random.randint(0, 1).dtype, mx.random.randint(0, 1, dtype=None).dtype ) + def test_randint_exact_integer_range(self): + # Bounds that are not exactly representable in float32 must not make + # values unreachable, return `high`, or collapse the interval. + key = mx.random.key(42) + + # At 2**24 adjacent float32 values are two integers apart. + a = mx.random.randint(2**24, 2**24 + 2, (10_000,), dtype=mx.int32, key=key) + values = set(a.tolist()) + self.assertNotIn(2**24 + 2, values) # `high` is exclusive + self.assertEqual(values, {2**24, 2**24 + 1}) # both are reachable + + # At 2**40 the float32 spacing is 2**17, which collapsed this interval + # to a single value. + b = mx.random.randint(2**40, 2**40 + 1024, (20_000,), dtype=mx.int64, key=key) + self.assertTrue(mx.all(b >= 2**40).item()) + self.assertTrue(mx.all(b < 2**40 + 1024).item()) + self.assertEqual(len(set(b.tolist())), 1024) + + # A width above 2**63 is representable only as unsigned. + c = mx.random.randint(-(2**62), 2**62, (10_000,), dtype=mx.int64, key=key) + self.assertTrue(mx.all(c >= -(2**62)).item()) + self.assertTrue(mx.all(c < 2**62).item()) + self.assertGreater(len(set(c.tolist())), 1) + + def test_randint_uniform_coverage(self): + # Every value in a small interval should be reachable, with roughly + # equal probability. + n = 200_000 + for low, high in [(0, 6), (-10, 10), (3, 15)]: + a = mx.random.randint(low, high, (n,), dtype=mx.int64) + counts = [mx.sum(a == v).item() for v in range(low, high)] + self.assertTrue(all(c > 0 for c in counts)) + expected = n / (high - low) + self.assertLess(max(abs(c - expected) for c in counts) / expected, 0.1) + def test_bernoulli(self): a = mx.random.bernoulli() self.assertEqual(a.shape, ()) diff --git a/tests/random_tests.cpp b/tests/random_tests.cpp index 04a86352f1..92223f6623 100644 --- a/tests/random_tests.cpp +++ b/tests/random_tests.cpp @@ -516,6 +516,36 @@ TEST_CASE("test random randint") { auto key = array({0, 0}, {1, 2}); CHECK_THROWS_AS( random::randint(low, high, {}, float32, key), std::invalid_argument); + + // Bounds that are not exactly representable in float32 must not make values + // unreachable, return `high`, or collapse the interval. + + // At 2^24 adjacent float32 values are two integers apart. + auto lo24 = array(1 << 24, int32); + auto hi24 = array((1 << 24) + 2, int32); + x = random::randint(lo24, hi24, {10000}, int32); + CHECK(all(less(x, hi24)).item()); // `high` is exclusive + CHECK(any(equal(x, lo24)).item()); + CHECK(any(equal(x, array((1 << 24) + 1, int32))).item()); + + // At 2^40 the float32 spacing is 2^17, which collapsed this interval to a + // single value. + auto lo40 = array(int64_t(1) << 40, int64); + auto hi40 = array((int64_t(1) << 40) + 1024, int64); + x = random::randint(lo40, hi40, {20000}, int64); + CHECK( + (all(less_equal(lo40, x)).item() && + all(less(x, hi40)).item())); + CHECK_GT(max(x).item(), min(x).item()); + + // A width above 2^63 is representable only as unsigned. + auto lo62 = array(-(int64_t(1) << 62), int64); + auto hi62 = array(int64_t(1) << 62, int64); + x = random::randint(lo62, hi62, {10000}, int64); + CHECK( + (all(less_equal(lo62, x)).item() && + all(less(x, hi62)).item())); + CHECK_GT(max(x).item(), min(x).item()); } TEST_CASE("test random bernoulli") {