Core Data Structures & Communication Primitives for Tensor Parallel for Keras

keras/api/_tf_keras/keras/distribution/__init__.py

@@ -15,6 +15,9 @@
     distribute_tensor as distribute_tensor,
 )
 from keras.src.distribution.distribution_lib import distribution as distribution
+from keras.src.distribution.distribution_lib import (
+    get_device_count as get_device_count,
+)
 from keras.src.distribution.distribution_lib import initialize as initialize
 from keras.src.distribution.distribution_lib import list_devices as list_devices
 from keras.src.distribution.distribution_lib import (

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/distribution/__init__.py

@@ -15,6 +15,9 @@
     distribute_tensor as distribute_tensor,
 )
 from keras.src.distribution.distribution_lib import distribution as distribution
+from keras.src.distribution.distribution_lib import (
+    get_device_count as get_device_count,
+)
 from keras.src.distribution.distribution_lib import initialize as initialize
 from keras.src.distribution.distribution_lib import list_devices as list_devices
 from keras.src.distribution.distribution_lib import (

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/core.py

@@ -1,5 +1,6 @@
 import jax
 import jax.experimental.sparse as jax_sparse
+import jax.lax as lax
 import jax.numpy as jnp
 import ml_dtypes
 import numpy as np
@@ -529,6 +530,61 @@ def remat(f):
     return jax.checkpoint(f)
 
 
+def all_reduce(x, op="sum", axis_name="model"):
+    """
+    Performs an **all-reduce** operation across all replicas in the specified
+    distribution axis.
+
+    The all-reduce operation computes a reduction (like sum, mean, or product)
+    of the input tensor `x` across all devices/replicas in the `axis_name`
+    group, and then broadcasts the result back to all participating devices.
+
+    Args:
+        x: The tensor to reduce.
+        op: The reduction operation to perform. Common options include "sum",
+            "mean", or "product". Defaults to "sum".
+        axis_name: The name of the distribution axis (e.g., "model",
+            "data") over which to perform the reduction. Defaults to "model".
+
+    Returns:
+        The result of the all-reduce operation, with the same shape as the
+        input `x`.
+    """
+    if op == "sum":
+        return lax.psum(x, axis_name=axis_name)
+    elif op == "mean":
+        return lax.pmean(x, axis_name=axis_name)
+    else:
+        raise ValueError(
+            f"Unsupported reduction operation: {op}. "
+            "Supported options are 'sum' and 'mean'."
+        )
+
+
+def all_gather(x, axis, axis_name="model"):
+    """
+    Performs an all-gather operation across all replicas in the specified
+    distribution axis.
+
+    The all-gather operation collects the input tensor `x` from all devices
+    in the `axis_name` group and concatenates them along the specified `axis`.
+    This is often used in tensor parallelism to combine parts of a tensor
+    distributed across devices.
+
+    Args:
+        x: The tensor to gather.
+        axis: The dimension along which to concatenate the gathered tensors.
+        axis_name: The name of the distribution axis (e.g., "model",
+                                   "data") over which to perform the gather.
+                                   Defaults to "model".
+
+    Returns:
+        The gathered tensor, which will have a larger size along `axis`
+        dimension.
+    """
+    return lax.all_gather(x, axis_name=axis_name, axis=axis, tiled=True)
+
+
 class name_scope(base_name_scope):
     def __init__(self, name, **kwargs):
         super().__init__(name, **kwargs)

keras/src/backend/jax/core_test.py

@@ -1,3 +1,4 @@
+import functools
 import os
 
 import jax
@@ -9,6 +10,8 @@
 from keras.src import backend
 from keras.src import testing
 from keras.src.backend.config import is_nnx_enabled
+from keras.src.backend.jax.core import all_gather
+from keras.src.backend.jax.core import all_reduce
 
 if is_nnx_enabled():
     from flax import nnx
@@ -66,3 +69,78 @@ def test_keras_variable_nnx_split_merge_sync(self):
         state = jax.tree.map(lambda x: x + 1, state)
         variable2 = nnx.merge(graphdef, state)
         self.assertEqual(variable2._value, variable2.value)
+
+
+@pytest.mark.skipif(
+    backend.backend() != "jax",
+    reason="JAX backend specific test for collective operations.",
+)
+@pytest.mark.skipif(
+    jax.local_device_count() < 2,
+    reason="Requires multiple local devices for testing.",
+)
+class JaxCollectiveOpsTest(testing.TestCase):
+    def test_all_reduce_sum(self):
+        """Tests the all_reduce operation with the 'sum' reduction."""
+        num_devices = jax.local_device_count()
+        local_value = 10.0
+
+        local_inputs = jax.numpy.array([local_value] * num_devices)
+
+        @functools.partial(
+            jax.pmap, axis_name="all", devices=jax.devices("cpu")
+        )
+        def reduce_sum_fn(x):
+            return all_reduce(x, op="sum", axis_name="all")
+
+        result = reduce_sum_fn(local_inputs)
+        expected_sum = local_value * num_devices
+
+        self.assertTrue(np.allclose(result, expected_sum))
+        self.assertEqual(result.shape, (num_devices,))
+
+    def test_all_reduce_mean(self):
+        """Tests the all_reduce operation with the 'mean' reduction."""
+        num_devices = jax.local_device_count()
+        local_value = 10.0
+
+        local_inputs = jax.numpy.array([local_value] * num_devices)
+
+        @functools.partial(
+            jax.pmap, axis_name="all", devices=jax.devices("cpu")
+        )
+        def reduce_mean_fn(x):
+            return all_reduce(x, op="mean", axis_name="all")
+
+        result = reduce_mean_fn(local_inputs)
+        expected_mean = local_value
+
+        self.assertTrue(np.allclose(result, expected_mean))
+        self.assertEqual(result.shape, (num_devices,))
+
+    def test_all_gather(self):
+        """Tests the all_gather operation."""
+        num_devices = jax.local_device_count()
+        local_data = np.arange(5)
+
+        local_inputs = jax.numpy.stack(
+            [local_data + (i * 5) for i in range(num_devices)]
+        )
+
+        @functools.partial(
+            jax.pmap, axis_name="all", devices=jax.devices("cpu")
+        )
+        def gather_fn(x):
+            return all_gather(x, axis=0, axis_name="all")
+
+        result_array_on_devices = gather_fn(local_inputs)
+
+        expected_shape = (num_devices, num_devices * local_data.shape[0])
+        self.assertEqual(result_array_on_devices.shape, expected_shape)
+
+        expected_gathered_data = np.arange(num_devices * local_data.shape[0])
+
+        for i in range(num_devices):
+            self.assertTrue(
+                np.allclose(result_array_on_devices[i], expected_gathered_data)
+            )

keras/src/backend/jax/distribution_lib.py

@@ -27,6 +27,21 @@ def list_devices(device_type=None):
     return [f"{device.platform}:{device.id}" for device in jax_devices]
 
 
+def get_device_count(device_type=None):
+    """Returns the number of available JAX devices.
+
+    Args:
+        device_type: Optional device type to count (e.g., "cpu", "gpu", "tpu").
+            If `None`, it counts all available devices.
+
+    Returns:
+        int: The total number of JAX devices for the specified type.
+    """
+    device_type = device_type.lower() if device_type else None
+    jax_devices = jax.devices(backend=device_type)
+    return len(jax_devices)
+
+
 def distribute_variable(value, layout):
     """Create a distributed variable for JAX.
 

keras/src/backend/jax/distribution_lib_test.py

@@ -29,8 +29,8 @@
 
 
 @pytest.mark.skipif(
-    backend.backend() != "jax",
-    reason="Backend specific test",
+    backend.backend() != "jax" or len(jax.devices()) != 8,
+    reason="Backend specific test and requires 8 devices",
 )
 class JaxDistributionLibTest(testing.TestCase):
     def _create_jax_layout(self, sharding):

keras/src/backend/jax/numpy.py

@@ -1167,6 +1167,10 @@ 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):
+    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,73 @@ def split(x, indices_or_sections, axis=0):
     )
 
 
+def array_split(x, indices_or_sections, axis=0):
+    x = get_ov_output(x)
+
+    if not isinstance(indices_or_sections, int):
+        raise TypeError(
+            "Argument `indices_or_sections` must be of type `int`. "
+            f"Received: {indices_or_sections}"
+        )
+    if indices_or_sections <= 0:
+        raise ValueError(
+            "Argument `indices_or_sections` must be a positive integer. "
+            f"Received: {indices_or_sections}"
+        )
+
+    num_splits_val = indices_or_sections
+    num_splits = ov_opset.constant(
+        np.array(num_splits_val, dtype=np.int64)
+    ).output(0)
+
+    axis_tensor = ov_opset.constant(np.array(axis, dtype=np.int64)).output(0)
+
+    zero_scalar = ov_opset.constant(np.array(0, dtype=np.int64)).output(0)
+
+    one_scalar = ov_opset.constant(np.array(1, dtype=np.int64)).output(0)
+
+    shape_tensor = ov_opset.shape_of(x, Type.i64).output(0)
+    axis_i64_vec = ov_opset.constant([axis], dtype=Type.i64).output(0)
+
+    total_size_tensor_vec = ov_opset.gather(
+        shape_tensor, axis_i64_vec, zero_scalar
+    ).output(0)
+
+    total_size = ov_opset.squeeze(total_size_tensor_vec, zero_scalar).output(0)
+
+    split_size = ov_opset.divide(
+        total_size, num_splits, auto_broadcast="NUMPY"
+    ).output(0)
+
+    remainder = ov_opset.mod(
+        total_size, num_splits, auto_broadcast="NUMPY"
+    ).output(0)
+
+    splits_shape = ov_opset.constant([num_splits_val], dtype=Type.i64).output(0)
+    all_splits_base = ov_opset.broadcast(split_size, splits_shape).output(0)
+
+    range_splits = ov_opset.range(
+        zero_scalar,
+        num_splits,
+        one_scalar,
+        Type.i64,
+    ).output(0)
+
+    remainder_bcast = ov_opset.broadcast(remainder, splits_shape).output(0)
+
+    add_one_mask = ov_opset.less(range_splits, remainder_bcast).output(0)
+
+    add_one_values = ov_opset.convert(add_one_mask, Type.i64).output(0)
+
+    split_lengths = ov_opset.add(all_splits_base, add_one_values).output(0)
+    splits = ov_opset.variadic_split(x, axis_tensor, split_lengths)
+
+    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,24 @@ 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 = tf.shape(x)[axis]
+    avg_size = tf.math.floordiv(total_size, num_splits)
+    remainder = tf.math.floormod(total_size, num_splits)
+
+    sizes = tf.concat(
+        [
+            tf.fill([remainder], avg_size + 1),
+            tf.fill([num_splits - remainder], avg_size),
+        ],
+        axis=0,
+    )
+
+    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,13 @@ 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)
+    axis_int = int(axis)
+    out = torch.tensor_split(x, indices_or_sections, dim=axis_int)
+    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/distribution/__init__.py

@@ -6,6 +6,7 @@
 from keras.src.distribution.distribution_lib import TensorLayout
 from keras.src.distribution.distribution_lib import distribute_tensor
 from keras.src.distribution.distribution_lib import distribution
+from keras.src.distribution.distribution_lib import get_device_count
 from keras.src.distribution.distribution_lib import initialize
 from keras.src.distribution.distribution_lib import list_devices
 from keras.src.distribution.distribution_lib import set_distribution