| 372 | self.assertEqual(sample.dtype, mx.complex64) |
| 373 | |
| 374 | def test_broadcastable_scale_loc(self): |
| 375 | b = mx.random.normal((10, 2)) |
| 376 | sample = mx.random.normal((2, 10, 2), loc=b, scale=b) |
| 377 | mx.eval(sample) |
| 378 | self.assertEqual(sample.shape, (2, 10, 2)) |
| 379 | |
| 380 | with self.assertRaises(ValueError): |
| 381 | b = mx.random.normal((10,)) |
| 382 | sample = mx.random.normal((2, 10, 2), loc=b, scale=b) |
| 383 | |
| 384 | b = mx.random.normal((3, 1, 2)) |
| 385 | sample = mx.random.normal((3, 4, 2), dtype=mx.float16, loc=b, scale=b) |
| 386 | mx.eval(sample) |
| 387 | self.assertEqual(sample.shape, (3, 4, 2)) |
| 388 | self.assertEqual(sample.dtype, mx.float16) |
| 389 | |
| 390 | |
| 391 | if __name__ == "__main__": |