zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Cast loc and scale to requested dtype

Browse Source
This commit is contained in:
Angelos Katharopoulos
2025-05-13 21:13:44 -07:00
parent 488537e80b
commit 4a380d5a80
2 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/tests/test_random.py
View File

@@ -381,6 +381,12 @@ class TestRandom(mlx_tests.MLXTestCase):
b = mx.random.normal((10,))
sample = mx.random.normal((2, 10, 2), loc=b, scale=b)
b = mx.random.normal((3, 1, 2))
sample = mx.random.normal((3, 4, 2), dtype=mx.float16, loc=b, scale=b)
mx.eval(sample)
self.assertEqual(sample.shape, (3, 4, 2))
self.assertEqual(sample.dtype, mx.float16)
if __name__ == "__main__":
unittest.main()