Fix RandomBits::is_equivalent to include width (#2978)

Co-authored-by: KD2YCU <[email protected]>
Co-authored-by: Angelos Katharopoulos <[email protected]>
Co-authored-by: Awni Hannun <[email protected]>

mlx/primitives.cpp

@@ -3717,7 +3717,7 @@ std::pair<std::vector<array>, std::vector<int>> RandomBits::vmap(
 
 bool RandomBits::is_equivalent(const Primitive& other) const {
   const RandomBits& r_other = static_cast<const RandomBits&>(other);
-  return shape_ == r_other.shape_;
+  return shape_ == r_other.shape_ && width_ == r_other.width_;
 }
 
 std::vector<array> Real::vjp(

python/tests/test_blas.py

@@ -297,6 +297,7 @@ class TestBlas(mlx_tests.MLXTestCase):
             self.assertTrue(np.allclose(out_mlx, out_npy, atol=1e-5))
 
     def test_matrix_vector(self):
+        mx.random.seed(0)
         for dtype in self.dtypes:
             with self.subTest(dtype=dtype):
                 np_dtype = getattr(np, dtype)

tests/compile_tests.cpp

@@ -803,3 +803,16 @@ TEST_CASE("test compile with no-ops") {
   auto out = compile(fun)({in})[0];
   CHECK_EQ(out.inputs()[0].id(), in.id());
 }
+
+TEST_CASE("test compile random bits") {
+  auto fun = [](const std::vector<array>& inputs) {
+    auto key = inputs[0];
+    auto a = random::bits({32, 32}, 4, key);
+    auto b = random::bits({32, 32}, 2, key);
+    return std::vector<array>{a + b};
+  };
+  auto in = random::key(0);
+  auto expected = fun({in})[0];
+  auto out = compile(fun)({in})[0];
+  CHECK(array_equal(out, expected).item<bool>());
+}