Sample randint directly instead of through a float32 uniform

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(

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, ())

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<bool>()); // `high` is exclusive
+  CHECK(any(equal(x, lo24)).item<bool>());
+  CHECK(any(equal(x, array((1 << 24) + 1, int32))).item<bool>());
+
+  // 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<bool>() &&
+       all(less(x, hi40)).item<bool>()));
+  CHECK_GT(max(x).item<int64_t>(), min(x).item<int64_t>());
+
+  // 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<bool>() &&
+       all(less(x, hi62)).item<bool>()));
+  CHECK_GT(max(x).item<int64_t>(), min(x).item<int64_t>());
 }
 
 TEST_CASE("test random bernoulli") {