code: def distribute_data_input(inputs, layout):
      """Distribute the input data with the corresponding layout.
  
      Note that the inputs here is a local worker batch. Within the local worker,
      the data need to be further partitioned to map to the each of the devices.
  
      Args:
          inputs: `jax.Array` that is already sharded to a local process size.
          layout: `TensorLayout` for the distribution information, or a
              `jax.sharding.Sharding` instance.
  
      Returns:
          Distributed inputs thats been properly put to local devices.
      """
      if not isinstance(layout, jax.sharding.Sharding):
          layout = _to_jax_layout(layout)
      if layout.is_fully_addressable:
          return jax.device_put(inputs, layout)
  
      # We need the jax mesh information to determine how to place the data
      # on to each of the worker.
      jax_mesh = layout.mesh
      mesh_rank = len(jax_mesh.shape)
      per_process_batch_size = inputs.shape[0]
      if mesh_rank == 1:
          # This is data parallel mesh only. We will split the full data
          # across the batch dim.
          num_split = jax.local_device_count()
          per_replica_batch_size = per_process_batch_size // num_split
          if per_process_batch_size % per_replica_batch_size != 0:
              raise ValueError(
                  f"The local batch size {per_process_batch_size} is not"
                  "divisible by the number of local replicas "
                  f"{num_split}"
              )
          global_batch_size = (
              per_process_batch_size * jax.process_count()
          )
          per_replica_batches = np.split(inputs, num_split, axis=0)
      elif mesh_rank == 2:
          # Data+Model parallel
          # In this case, we need to check if the mesh batch dim shape is large
          # than number of local devices, so that we can decide whether a split
          # is needed for the data, or a repeat/copy of the data is needed for
          # each of the device.
          # TODO(scottzhu): The mesh batch dim name is not available here, since
          # we only have jax Mesh. We assume the first dim is for batch, and
          # second dim is for model for now.
          mesh_batch_dim_size = list(jax_mesh.shape.values())[0]
          local_device_count = jax.local_device_count()
          if mesh_batch_dim_size < local_device_count:
              # No split needed, we only need to repeat here.
              global_batch_size = per_process_batch_size
              per_replica_batches = [
                  inputs for _ in range(local_device_count)
              ]
          else:
              # Note that global batch size is not simply per_process_batch_size *
              # num_process. It actually depends on the model dim size.
              global_batch_size = per_process_batch_size * (
                  mesh_batch_dim_size // local_device_count
              )
              per_replica_batches = np.split(
                  inputs, local_device_count, axis=0
              )
      else:
          raise ValueError(
              "Only 1D or 2D mesh is supported at the moment. "
              f"Received mesh shape = {jax_mesh.shape}"
          )
  
      global_shape = (global_batch_size,) + inputs.shape[1:]
      global_batch_array = jax.make_array_from_single_device_arrays(
          global_shape,
          layout,
          arrays=[
              jax.device_put(batch, device)
              for batch, device in zip(
                  per_replica_batches, layout.addressable_devices
              )
          ],
      )
      return global_batch_array
  
file_path: keras/keras/backend/jax/distribution_lib.py
imports: import jax
  import numpy
  from keras.utils import jax_utils
label_desc: Use image node_id 4 for calling a global function with example usage: distribute_data_input(inputs, layout) and returns: global_batch_array, jax