| 258 | mx.random.bernoulli(0, [2]) # Bad type |
| 259 | |
| 260 | def test_truncated_normal(self): |
| 261 | a = mx.random.truncated_normal(-2.0, 2.0) |
| 262 | self.assertEqual(a.size, 1) |
| 263 | self.assertEqual(a.dtype, mx.float32) |
| 264 | |
| 265 | a = mx.random.truncated_normal(mx.array([]), mx.array([])) |
| 266 | self.assertEqual(a.dtype, mx.float32) |
| 267 | self.assertEqual(a.size, 0) |
| 268 | |
| 269 | lower = mx.reshape(mx.array([-2.0, 0.0]), [1, 2]) |
| 270 | upper = mx.reshape(mx.array([0.0, 1.0, 2.0]), [3, 1]) |
| 271 | a = mx.random.truncated_normal(lower, upper) |
| 272 | |
| 273 | self.assertEqual(a.shape, (3, 2)) |
| 274 | self.assertTrue(mx.all(lower <= a).item() and mx.all(a <= upper).item()) |
| 275 | |
| 276 | a = mx.random.truncated_normal(2.0, -2.0) |
| 277 | self.assertTrue(mx.all(a == 2.0).item()) |
| 278 | |
| 279 | a = mx.random.truncated_normal(-3.0, 3.0, [542, 399]) |
| 280 | self.assertEqual(a.shape, (542, 399)) |
| 281 | |
| 282 | lower = mx.array([-2.0, -1.0]) |
| 283 | higher = mx.array([1.0, 2.0, 3.0]) |
| 284 | with self.assertRaises(ValueError): |
| 285 | mx.random.truncated_normal(lower, higher) # Bad shape |
| 286 | |
| 287 | self.assertEqual( |
| 288 | mx.random.truncated_normal(0, 1).dtype, |
| 289 | mx.random.truncated_normal(0, 1, dtype=None).dtype, |
| 290 | ) |
| 291 | |
| 292 | def test_gumbel(self): |
| 293 | samples = mx.random.gumbel(shape=(100, 100)) |