From a25399cbd48b80052bbfc6495ca5a195f9b71aab Mon Sep 17 00:00:00 2001
From: Dan Anderson <dan@kd2ycu.com>
Date: Tue, 10 Mar 2026 00:47:08 -0400
Subject: [PATCH] Validate num_splits in split (#3234)

Co-authored-by: KD2YCU <me@kd2ycu.com>
---
 mlx/ops.cpp              |  8 +++++++-
 python/tests/test_ops.py | 19 +++++++++++++++++++
 tests/ops_tests.cpp      |  4 ++++
 3 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/mlx/ops.cpp b/mlx/ops.cpp
index c7af8834..e6554c2c 100644
--- a/mlx/ops.cpp
+++ b/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
\ No newline at end of file
+} // namespace mlx::core
diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py
index fd40f3b6..790689e7 100644
--- a/python/tests/test_ops.py
+++ b/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)
diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp
index 62fd8c59..6a924cfb 100644
--- a/tests/ops_tests.cpp
+++ b/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);