| 303 | ) |
| 304 | |
| 305 | def test_categorical(self): |
| 306 | logits = mx.zeros((10, 20)) |
| 307 | self.assertEqual(mx.random.categorical(logits, -1).shape, (10,)) |
| 308 | self.assertEqual(mx.random.categorical(logits, 0).shape, (20,)) |
| 309 | self.assertEqual(mx.random.categorical(logits, 1).shape, (10,)) |
| 310 | |
| 311 | out = mx.random.categorical(logits) |
| 312 | self.assertEqual(out.shape, (10,)) |
| 313 | self.assertEqual(out.dtype, mx.uint32) |
| 314 | self.assertTrue(mx.max(out).item() < 20) |
| 315 | |
| 316 | out = mx.random.categorical(logits, 0, [5, 20]) |
| 317 | self.assertEqual(out.shape, (5, 20)) |
| 318 | self.assertTrue(mx.max(out).item() < 10) |
| 319 | |
| 320 | out = mx.random.categorical(logits, 1, num_samples=7) |
| 321 | self.assertEqual(out.shape, (10, 7)) |
| 322 | out = mx.random.categorical(logits, 0, num_samples=7) |
| 323 | self.assertEqual(out.shape, (20, 7)) |
| 324 | |
| 325 | with self.assertRaises(ValueError): |
| 326 | mx.random.categorical(logits, shape=[10, 5], num_samples=5) |
| 327 | |
| 328 | def test_permutation(self): |
| 329 | x = sorted(mx.random.permutation(4).tolist()) |