code: def distribute_variable(value, layout):
      """Create a distributed variable for JAX.
  
      Since JAX doesn't have a variable class, this will just return a `jax.Array`
      with the corresponding layout/sharding specified.
  
      Note that this function should be used in eager context, not in jitted
      function.
  
      Args:
          value: the initial value of the variable.
          layout: `TensorLayout` for the created variable, or a
              `jax.sharding.Sharding` instance.
  
      Returns:
          jax.Array which is the distributed variable.
      """
      if not isinstance(layout, jax.sharding.Sharding):
          layout = _to_jax_layout(layout)
      if isinstance(
          value, (jax.Array, jax.numpy.ndarray)
      ) and value.sharding.is_equivalent_to(
          layout, ndim=len(value.shape)
      ):
          # Skip the relayout if the value is already having the proper sharding
          return value
  
      if layout.is_fully_addressable:
          return jax.device_put(value, layout)
      else:
          # Need to only distribute the value to local addressible devices, and
          # repack them back into global format.
          mapping = layout.addressable_devices_indices_map(value.shape)
          local_values = jax.device_put(
              [value[i] for i in mapping.values()], list(mapping.keys())
          )
          global_value = jax.make_array_from_single_device_arrays(
              value.shape, layout, local_values
          )
          return global_value
  
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 2 for calling a global function with example usage: distribute_variable(value, layout) and returns: value, jax, global_value