Repository: keras-team/keras-core (language: python)

Issue:
Behavior of ops.split
(This was originally found while working on [keras-nlp\#1110](https://github.com/keras-team/keras-nlp/pull/1110) )

An interesting revelation on the behavior of the split method across torch, tf and numpy

```
import torch
a = torch.ones(2, 8, 2, 0)
x = torch.split(a, 2, dim=-1)
print("splits = ", x)
```

output 
splits =  (tensor([], size=(2, 8, 2, 0)),)

```
import tensorflow as tf
b = tf.ones((2, 8, 2, 0))
y = tf.split(b, 2, axis=-1)
print("splits = ", y)
```

output 
splits =  [<tf.Tensor: shape=(2, 8, 2, 0), dtype=float32, numpy=array([], shape=(2, 8, 2, 0), dtype=float32)>, <tf.Tensor: shape=(2, 8, 2, 0), dtype=float32, numpy=array([], shape=(2, 8, 2, 0), dtype=float32)>]

```
import numpy as np

c = np.ones((2, 8, 2, 0))
z = np.split(c, 2, axis=-1)
print("splits = ", z)
```

output 
splits =  [array([], shape=(2, 8, 2, 0), dtype=float64), array([], shape=(2, 8, 2, 0), dtype=float64)]

How shall we handle this behavior to deliver same functionality !


Discussion / hints:
What about JAX?

It sounds like TF and numpy are consistent (apart from the fact that np uses float64, which is expected).

Our general approach is that we standardize on the numpy API. So here we should fix the torch op to make sure it follows the numpy behavior. Are you able to open a PR?
IIRC JAX follows the same convention as numpy. So, apart from torch everything else seems consistent. But do we know the reasoning behind the divergence in torch? @shivance 
Yes Jax has it's own numpy API which is consistent with numpy, regarding this behavior of Torch, I'll need to dig in. Currently pulled up in couple of PRs at KerasNLP end. Will followup in this thread or a PR by this weekend if that's okay ?
Sounds great -- thank you!

Fix patch:
diff --git a/keras_core/backend/torch/numpy.py b/keras_core/backend/torch/numpy.py
index ba613cfcd..4986320f6 100644
--- a/keras_core/backend/torch/numpy.py
+++ b/keras_core/backend/torch/numpy.py
@@ -800,22 +800,34 @@ def sort(x, axis=-1):
 
 def split(x, indices_or_sections, axis=0):
     x = convert_to_tensor(x)
+    dim = x.shape[axis]
     if isinstance(indices_or_sections, (list, tuple)):
         idxs = convert_to_tensor(indices_or_sections)
         start_size = indices_or_sections[0]
-        end_size = x.shape[axis] - indices_or_sections[-1]
+        end_size = dim - indices_or_sections[-1]
         chunk_sizes = (
             [start_size]
             + torch.diff(idxs).type(torch.int).tolist()
             + [end_size]
         )
     else:
-        chunk_sizes = x.shape[axis] // indices_or_sections
-    return torch.split(
+        if dim % indices_or_sections != 0:
+            raise ValueError(
+                f"Received indices_or_sections={indices_or_sections} "
+                f"(interpreted as a number of sections) and axis={axis}, "
+                f"but input dimension x.shape[{axis}]={x.shape[axis]} "
+                f"is not divisible by {indices_or_sections}. "
+                f"Full input shape: x.shape={x.shape}"
+            )
+        chunk_sizes = dim // indices_or_sections
+    out = torch.split(
         tensor=x,
         split_size_or_sections=chunk_sizes,
         dim=axis,
     )
+    if dim == 0 and isinstance(indices_or_sections, int):
+        out = tuple(out[0].clone() for _ in range(indices_or_sections))
+    return out
 
 
 def stack(x, axis=0):


Test patch:
diff --git a/keras_core/ops/numpy_test.py b/keras_core/ops/numpy_test.py
index 0ff0d364b..2ae7f12c6 100644
--- a/keras_core/ops/numpy_test.py
+++ b/keras_core/ops/numpy_test.py
@@ -3283,40 +3283,25 @@ def test_sort(self):
 
     def test_split(self):
         x = np.array([[1, 2, 3], [3, 2, 1]])
-        if backend.backend() == "torch":
-            self.assertAllClose(
-                [backend.convert_to_numpy(t) for t in knp.split(x, 2)],
-                np.split(x, 2),
-            )
-            self.assertAllClose(
-                [backend.convert_to_numpy(t) for t in knp.Split(2)(x)],
-                np.split(x, 2),
-            )
-            self.assertAllClose(
-                [
-                    backend.convert_to_numpy(t)
-                    for t in knp.split(x, [1, 2], axis=1)
-                ],
-                np.split(x, [1, 2], axis=1),
-            )
-            self.assertAllClose(
-                [
-                    backend.convert_to_numpy(t)
-                    for t in knp.Split([1, 2], axis=1)(x)
-                ],
-                np.split(x, [1, 2], axis=1),
-            )
-        else:
-            self.assertAllClose(knp.split(x, 2), np.split(x, 2))
-            self.assertAllClose(knp.Split(2)(x), np.split(x, 2))
-            self.assertAllClose(
-                knp.split(x, [1, 2], axis=1),
-                np.split(x, [1, 2], axis=1),
-            )
-            self.assertAllClose(
-                knp.Split([1, 2], axis=1)(x),
-                np.split(x, [1, 2], axis=1),
-            )
+        self.assertAllClose(knp.split(x, 2), np.split(x, 2))
+        self.assertAllClose(knp.Split(2)(x), np.split(x, 2))
+        self.assertAllClose(
+            knp.split(x, [1, 2], axis=1),
+            np.split(x, [1, 2], axis=1),
+        )
+        self.assertAllClose(
+            knp.Split([1, 2], axis=1)(x),
+            np.split(x, [1, 2], axis=1),
+        )
+
+        # test invalid indices_or_sections
+        with self.assertRaises(Exception):
+            knp.split(x, 3)
+
+        # test zero dimension
+        x = np.ones(shape=(0,))
+        self.assertEqual(len(knp.split(x, 2)), 2)
+        self.assertEqual(len(knp.Split(2)(x)), 2)
 
     def test_sqrt(self):
         x = np.array([[1, 4, 9], [16, 25, 36]])