Add keras.ops.array_split for Tensor Parallelism Support (#21697)

* Added tensor parallel for keras (Part 1/3)

* Removed unnecessary lines

* Fixes suggested by Gemini

* Fixes suggested by Gemini

* Fixes suggested by Gemini

* Fixes suggested by Gemini

* Fixes suggested by Gemini

* Fixes suggested by Gemini

* Fixing the failing test

* Fixing the failing test

* Fixing test

* Adding tests for distributed_backends

* Modifications for failing tests

* Modified for failing test

* Modified for failing test

* Modified for failing test

* added debuggers

* removed debuggers

* Removed the tensorflow, numpy and torch backends

* Refactoring the code

* Refactoring the code

* refactoring

* Adding necessary docstrings

* Removing redundancies

* Modifying tests

* Reformatting

* Reformatting the code

* Fixing failing tests

* fixes

* Fixing tests

* formatting

* fixing test

* fixing test

* Removing redundant lines

* Refactoring to remove communications.py and state_action_keras.py

* formatting the files

* fixing skip issues

* fixing test

* fixing test

* refactoring to remove distributed backend wrapper

* fixing test

* making distrubed backend more jax friendly

* Fixing comments

* Fixing comments

* Fixing comments

* fixes

* Refactor

* refactoring to resolve comments

* fixes

* fixes

* fix

* fix

* removing get_best_devices

* fixing comments

* fixing merge conflict

* modifying variable name

* Fixes

* fix

* fix

* splitting into 3 PRs

* Modified array_split implementation in openvino, tensorflow and torch

* formatting the array split function

* adding test for uneven array split

* fixing test

Author: buildwithsuhana

keras/api/_tf_keras/keras/ops/__init__.py

@@ -140,6 +140,7 @@
 from keras.src.ops.numpy import argpartition as argpartition
 from keras.src.ops.numpy import argsort as argsort
 from keras.src.ops.numpy import array as array
+from keras.src.ops.numpy import array_split as array_split
 from keras.src.ops.numpy import average as average
 from keras.src.ops.numpy import bartlett as bartlett
 from keras.src.ops.numpy import bincount as bincount

keras/api/_tf_keras/keras/ops/numpy/__init__.py

@@ -26,6 +26,7 @@
 from keras.src.ops.numpy import argpartition as argpartition
 from keras.src.ops.numpy import argsort as argsort
 from keras.src.ops.numpy import array as array
+from keras.src.ops.numpy import array_split as array_split
 from keras.src.ops.numpy import average as average
 from keras.src.ops.numpy import bartlett as bartlett
 from keras.src.ops.numpy import bincount as bincount

keras/api/ops/__init__.py

@@ -140,6 +140,7 @@
 from keras.src.ops.numpy import argpartition as argpartition
 from keras.src.ops.numpy import argsort as argsort
 from keras.src.ops.numpy import array as array
+from keras.src.ops.numpy import array_split as array_split
 from keras.src.ops.numpy import average as average
 from keras.src.ops.numpy import bartlett as bartlett
 from keras.src.ops.numpy import bincount as bincount

keras/api/ops/numpy/__init__.py

@@ -26,6 +26,7 @@
 from keras.src.ops.numpy import argpartition as argpartition
 from keras.src.ops.numpy import argsort as argsort
 from keras.src.ops.numpy import array as array
+from keras.src.ops.numpy import array_split as array_split
 from keras.src.ops.numpy import average as average
 from keras.src.ops.numpy import bartlett as bartlett
 from keras.src.ops.numpy import bincount as bincount

keras/src/backend/jax/numpy.py

@@ -1167,6 +1167,11 @@ def split(x, indices_or_sections, axis=0):
     return jnp.split(x, indices_or_sections, axis=axis)
 
 
+def array_split(x, indices_or_sections, axis=0):
+    x = convert_to_tensor(x)
+    return jnp.array_split(x, indices_or_sections, axis=axis)
+
+
 def stack(x, axis=0):
     x = [convert_to_tensor(t) for t in x]
     return jnp.stack(x, axis=axis)

keras/src/backend/numpy/numpy.py

@@ -1107,6 +1107,11 @@ def split(x, indices_or_sections, axis=0):
     return np.split(x, indices_or_sections, axis=axis)
 
 
+def array_split(x, indices_or_sections, axis=0):
+    axis = standardize_axis_for_numpy(axis)
+    return np.array_split(x, indices_or_sections, axis=axis)
+
+
 def stack(x, axis=0):
     axis = standardize_axis_for_numpy(axis)
     dtype_set = set([getattr(a, "dtype", type(a)) for a in x])

keras/src/backend/openvino/numpy.py

@@ -2014,6 +2014,37 @@ def split(x, indices_or_sections, axis=0):
     )
 
 
+def array_split(x, indices_or_sections, axis=0):
+    original_shape = x.shape
+    x = get_ov_output(x)
+
+    num_splits_val = indices_or_sections
+    total_size = original_shape[axis]
+    if total_size is None:
+        raise ValueError(
+            f"Cannot use array_split with static Python logic on dynamic axis. "
+            f"Axis {axis} has unknown dimension for shape {original_shape}."
+        )
+
+    base_size = total_size // num_splits_val
+    remainder = total_size % num_splits_val
+
+    split_lengths = [base_size + 1] * remainder + [base_size] * (
+        num_splits_val - remainder
+    )
+    split_lengths_tensor = ov_opset.constant(
+        split_lengths, dtype=Type.i64
+    ).output(0)
+
+    axis_tensor = ov_opset.constant(axis, dtype=Type.i32).output(0)
+    splits = ov_opset.variadic_split(x, axis_tensor, split_lengths_tensor)
+
+    result = []
+    for i in range(num_splits_val):
+        result.append(OpenVINOKerasTensor(splits.output(i)))
+    return result
+
+
 def stack(x, axis=0):
     if isinstance(x, tuple):
         x = list(x)

keras/src/backend/tensorflow/numpy.py

@@ -2494,6 +2494,17 @@ def split(x, indices_or_sections, axis=0):
     return tf.split(x, num_or_size_splits, axis=axis)
 
 
+def array_split(x, indices_or_sections, axis=0):
+    x = tf.convert_to_tensor(x)
+    num_splits = indices_or_sections
+    total_size = shape_op(x)[axis]
+    avg_size = total_size // num_splits
+    remainder = total_size % num_splits
+    sizes = [avg_size + 1] * remainder + [avg_size] * (num_splits - remainder)
+
+    return tf.split(x, sizes, axis=axis)
+
+
 def stack(x, axis=0):
     dtype_set = set([getattr(a, "dtype", type(a)) for a in x])
     if len(dtype_set) > 1:

keras/src/backend/torch/numpy.py

@@ -1539,6 +1539,12 @@ def split(x, indices_or_sections, axis=0):
     return list(out)
 
 
+def array_split(x, indices_or_sections, axis=0):
+    x = convert_to_tensor(x)
+    out = torch.tensor_split(x, indices_or_sections, dim=axis)
+    return list(out)
+
+
 def stack(x, axis=0):
     x = [convert_to_tensor(elem) for elem in x]
     return torch.stack(x, dim=axis)

keras/src/ops/numpy.py

@@ -7762,3 +7762,104 @@ def histogram(x, bins=10, range=None):
             f"Received: input.shape={x.shape}"
         )
     return backend.numpy.histogram(x, bins=bins, range=range)
+
+
+class ArraySplit(Operation):
+    def __init__(self, indices_or_sections, axis=0, *, name=None):
+        super().__init__(name=name)
+
+        self.indices_or_sections = indices_or_sections
+        self.axis = axis
+
+    def call(self, x):
+        return backend.numpy.array_split(
+            x,
+            indices_or_sections=self.indices_or_sections,
+            axis=self.axis,
+        )
+
+    def compute_output_spec(self, x):
+        num_splits = self.indices_or_sections
+
+        axis = self.axis
+        if axis < 0:
+            axis += len(x.shape)
+
+        total_size = x.shape[axis]
+
+        if total_size is None:
+            output_specs = []
+            base_shape = list(x.shape)
+            base_shape[axis] = None
+            for _ in range(num_splits):
+                output_specs.append(
+                    KerasTensor(shape=tuple(base_shape), dtype=x.dtype)
+                )
+            return tuple(output_specs)
+
+        split_size = total_size // num_splits
+        remainder = total_size % num_splits
+
+        output_specs = []
+        base_shape = list(x.shape)
+        for i in range(num_splits):
+            size = split_size + (1 if i < remainder else 0)
+            shape = base_shape.copy()
+            shape[axis] = size
+            output_specs.append(KerasTensor(shape=tuple(shape), dtype=x.dtype))
+
+        return list(output_specs)
+
+
+@keras_export(["keras.ops.array_split", "keras.ops.numpy.array_split"])
+def array_split(x, indices_or_sections, axis=0):
+    """Splits an array into multiple sub-arrays (unevenly).
+
+    This is similar to `keras.ops.split`, but it allows for
+    unequal splits. `indices_or_sections` must be an integer
+    that indicates the total number of sub-arrays to create.
+    If the tensor cannot be divided evenly, the first `remainder`
+    splits will have size `quotient + 1`, and the rest will
+    have size `quotient`.
+
+    Args:
+        x: Input tensor.
+        indices_or_sections: An integer indicating the number of
+            sub-arrays to create.
+        axis: The axis along which to split. Defaults to 0.
+
+    Returns:
+        A list of sub-tensors.
+
+    Example:
+    >>> x = keras.ops.arange(10)
+    >>> keras.ops.array_split(x, 3)
+    (array([0, 1, 2, 3], dtype=int32),
+     array([4, 5, 6], dtype=int32),
+     array([7, 8, 9], dtype=int32))
+    """
+    if not isinstance(indices_or_sections, int):
+        raise TypeError(
+            "Argument `indices_or_sections` must be of type `int`. "
+            f"Received: indices_or_sections={indices_or_sections}"
+        )
+
+    if indices_or_sections <= 0:
+        raise ValueError(
+            "Argument `indices_or_sections` must be a positive integer. "
+            f"Received: indices_or_sections={indices_or_sections}"
+        )
+
+    if not isinstance(axis, int):
+        raise TypeError(
+            f"Argument `axis` must be of type `int`. Received: {axis}"
+        )
+
+    if any_symbolic_tensors((x,)):
+        return ArraySplit(
+            indices_or_sections=indices_or_sections, axis=axis
+        ).symbolic_call(x)
+
+    return backend.numpy.array_split(
+        x, indices_or_sections=indices_or_sections, axis=axis
+    )