code: def _to_jax_layout(tensor_layout):
      """Convert the TensorLayout to JAX backend specific Sharding.
  
      Args:
          tensor_layout: TensorLayout instance to convert.
  
      Returns:
          A `jax.sharding.NamedSharding` instance.
      """
      if tensor_layout.device_mesh is None:
          raise ValueError(
              "Cannot create sharding when device mesh is not set "
              "for TensorLayout."
          )
      partition_spec = jax.sharding.PartitionSpec(*tensor_layout.axes)
      jax_mesh = _to_jax_mesh(tensor_layout.device_mesh)
      return jax.sharding.NamedSharding(jax_mesh, partition_spec)
  
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 10 for calling a global function with example usage: _to_jax_layout(tensor_layout) and returns: jax