| 200 | check_jointly_gaussian(data, mean, cov) |
| 201 | |
| 202 | def test_randint(self): |
| 203 | a = mx.random.randint(0, 1, []) |
| 204 | self.assertEqual(a.shape, ()) |
| 205 | self.assertEqual(a.dtype, mx.int32) |
| 206 | |
| 207 | shape = (88,) |
| 208 | low = mx.array(3) |
| 209 | high = mx.array(15) |
| 210 | |
| 211 | key = mx.random.key(0) |
| 212 | a = mx.random.randint(low, high, shape, key=key) |
| 213 | self.assertEqual(a.shape, shape) |
| 214 | self.assertEqual(a.dtype, mx.int32) |
| 215 | |
| 216 | # Check using the same key yields the same value |
| 217 | b = mx.random.randint(low, high, shape, key=key) |
| 218 | self.assertListEqual(a.tolist(), b.tolist()) |
| 219 | |
| 220 | shape = (3, 4) |
| 221 | low = mx.reshape(mx.array([0] * 3), [3, 1]) |
| 222 | high = mx.reshape(mx.array([12, 13, 14, 15]), [1, 4]) |
| 223 | |
| 224 | a = mx.random.randint(low, high, shape) |
| 225 | self.assertEqual(a.shape, shape) |
| 226 | |
| 227 | a = mx.random.randint(-10, 10, [1000, 1000]) |
| 228 | self.assertTrue(mx.all(-10 <= a).item() and mx.all(a < 10).item()) |
| 229 | |
| 230 | a = mx.random.randint(10, -10, [1000, 1000]) |
| 231 | self.assertTrue(mx.all(a == 10).item()) |
| 232 | |
| 233 | self.assertEqual( |
| 234 | mx.random.randint(0, 1).dtype, mx.random.randint(0, 1, dtype=None).dtype |
| 235 | ) |
| 236 | |
| 237 | def test_bernoulli(self): |
| 238 | a = mx.random.bernoulli() |