zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 00:04:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Minor updates to address a few issues (#537)

Browse Source 
* docs on arg indices return type

* arange with nan

* undo isort
This commit is contained in:
Awni Hannun
2024-01-23 22:24:41 -08:00
committed by GitHub
parent 4fe2fa2a64
commit f30e63353a
3 changed files with 23 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/tests/test_ops.py
View File

@@ -980,6 +980,17 @@ class TestOps(mlx_tests.MLXTestCase):
self.assertEqual(z.tolist(), [5, 6, 7])
def test_arange_overload_dispatch(self):
with self.assertRaises(ValueError):
a = mx.arange(float("nan"), 1, 5)
with self.assertRaises(ValueError):
a = mx.arange(0, float("nan"), 5)
with self.assertRaises(ValueError):
a = mx.arange(0, 2, float("nan"))
with self.assertRaises(ValueError):
a = mx.arange(0, float("inf"), float("inf"))
with self.assertRaises(ValueError):
a = mx.arange(float("inf"), 1, float("inf"))
a = mx.arange(5)
expected = [0, 1, 2, 3, 4]
self.assertListEqual(a.tolist(), expected)