improved error msg for invalid axis(mx.split) (#685)

* improved error msg for invalid axis(`mx.split`) * Apply suggestions from code review Co-authored-by: Awni Hannun <[email protected]> * fixed formatting issue --------- Co-authored-by: Awni Hannun <[email protected]>
2024-02-15 07:25:38 -08:00
parent 35431a4ac8
commit 85143fecdd
2 changed files with 10 additions and 0 deletions
@@ -1027,6 +1027,9 @@ class TestOps(mlx_tests.MLXTestCase):
        self.assertEqual(y.tolist(), [[3, 4]])
        self.assertEqual(z.tolist(), [[5, 6]])

+        with self.assertRaises(ValueError):
+            mx.split(a, 3, axis=2)
+
        a = mx.arange(8)
        x, y, z = mx.split(a, [1, 5])
        self.assertEqual(x.tolist(), [0])