Validate num_splits in split (#3234)

Co-authored-by: KD2YCU <[email protected]>

mlx/ops.cpp

@@ -955,6 +955,12 @@ split(const array& a, int num_splits, int axis, StreamOrDevice s /* = {} */) {
         << " for array with shape " << a.shape() << ".";
     throw std::invalid_argument(msg.str());
   }
+  if (num_splits <= 0) {
+    std::ostringstream msg;
+    msg << "[split] num_splits must be positive and non-zero but got "
+        << num_splits << ".";
+    throw std::invalid_argument(msg.str());
+  }
   auto q_and_r = std::ldiv(a.shape(axis), num_splits);
   if (q_and_r.rem) {
     std::ostringstream msg;
@@ -6264,4 +6270,4 @@ array contiguous(
       {a});
 }
 
-} // namespace mlx::core
+} // namespace mlx::core

python/tests/test_ops.py

@@ -1328,6 +1328,25 @@ class TestOps(mlx_tests.MLXTestCase):
         self.assertEqual(y.tolist(), [1, 2, 3, 4])
         self.assertEqual(z.tolist(), [5, 6, 7])
 
+    def test_split_invalid_num_splits(self):
+        """Regression: split with num_splits <= 0 should raise, not crash."""
+        a = mx.arange(6)
+
+        # num_splits = 0: should raise cleanly (was UB via divide-by-zero)
+        with self.assertRaises(ValueError):
+            mx.split(a, 0)
+
+        # num_splits = -1: should raise cleanly (was SIGBUS via huge allocation)
+        with self.assertRaises(ValueError):
+            mx.split(a, -1)
+
+        # Also check with explicit axis
+        b = mx.zeros((4, 6))
+        with self.assertRaises(ValueError):
+            mx.split(b, 0, axis=1)
+        with self.assertRaises(ValueError):
+            mx.split(b, -2, axis=0)
+
     def test_arange_overload_dispatch(self):
         with self.assertRaises(ValueError):
             a = mx.arange(float("nan"), 1, 5)

tests/ops_tests.cpp

@@ -415,6 +415,10 @@ TEST_CASE("test split") {
   array x = array(1);
   CHECK_THROWS(split(x, 0));
 
+  // Regression: non-scalar split with num_splits <= 0
+  CHECK_THROWS(split(array({0, 1, 2, 3, 4, 5}), 0));
+  CHECK_THROWS(split(array({0, 1, 2, 3, 4, 5}), -1));
+
   x = array({3});
   CHECK_EQ(split(x, 1)[0].item<int>(), 3);