| 290 | ) |
| 291 | |
| 292 | def test_gumbel(self): |
| 293 | samples = mx.random.gumbel(shape=(100, 100)) |
| 294 | self.assertEqual(samples.shape, (100, 100)) |
| 295 | self.assertEqual(samples.dtype, mx.float32) |
| 296 | mean = 0.5772 |
| 297 | # Std deviation of the sample mean is small (<0.02), |
| 298 | # so this test is pretty conservative |
| 299 | self.assertTrue(mx.abs(mx.mean(samples) - mean) < 0.2) |
| 300 | |
| 301 | self.assertEqual( |
| 302 | mx.random.gumbel((1, 1)).dtype, mx.random.gumbel((1, 1), dtype=None).dtype |
| 303 | ) |
| 304 | |
| 305 | def test_categorical(self): |
| 306 | logits = mx.zeros((10, 20)) |