Remove references to jax.experimental.layout.Layout. (#21400)

It is being renamed to jax.experimental.layout.Format.  The comments
can be kept generic anyways to support other layout instances
in the future.

Author: cantonios

keras/src/backend/jax/distribution_lib.py

@@ -2,7 +2,6 @@
 
 import jax
 import numpy as np
-from jax.experimental import layout as jax_layout
 
 from keras.src.backend.common import global_state
 from keras.src.random import seed_generator
@@ -40,8 +39,7 @@ def distribute_variable(value, layout):
     Args:
         value: the initial value of the variable.
         layout: `TensorLayout` for the created variable, or a
-            JAX-supported layout instance
-            (e.g. `jax.experimental.layout.Layout`, `jax.sharding.Sharding`).
+            JAX-supported layout instance (e.g. `jax.sharding.Sharding`).
 
     Returns:
         jax.Array which is the distributed variable.
@@ -58,8 +56,7 @@ def distribute_tensor(tensor, layout):
     Args:
         tensor: `jax.Array` that need to be distributed.
         layout: `TensorLayout` for the created variable, or a
-            JAX-supported layout instance
-            (e.g. `jax.experimental.layout.Layout`, `jax.sharding.Sharding`).
+            JAX-supported layout instance (e.g. `jax.sharding.Sharding`).
 
     Returns:
         Distributed value.
@@ -81,7 +78,8 @@ def distribute_tensor(tensor, layout):
             layout, jax.sharding.Sharding
         ) and tensor.sharding.is_equivalent_to(layout, ndim=len(tensor.shape)):
             return tensor
-        elif isinstance(layout, jax_layout.Layout):
+        # JAX explicit layout support.
+        elif hasattr(layout, "layout"):
             current_layout = getattr(tensor, "layout", None)
             if current_layout == layout:
                 return tensor

keras/src/backend/jax/distribution_lib_test.py

@@ -33,6 +33,15 @@
     reason="Backend specific test",
 )
 class JaxDistributionLibTest(testing.TestCase):
+    def _create_jax_layout(self, sharding):
+        # Use jax_layout.Format or jax_layout.Layout if available.
+        if hasattr(jax_layout, "Format"):
+            return jax_layout.Format(sharding=sharding)
+        elif hasattr(jax_layout, "Layout"):
+            return jax_layout.Layout(sharding=sharding)
+
+        return sharding
+
     def test_list_devices(self):
         self.assertEqual(len(distribution_lib.list_devices()), 8)
         self.assertEqual(len(distribution_lib.list_devices("cpu")), 8)
@@ -132,7 +141,7 @@ def test_distribute_tensor_with_jax_layout(self):
         )
 
         inputs = jax.numpy.array(np.random.normal(size=(16, 8)))
-        target_layout = jax_layout.Layout(
+        target_layout = self._create_jax_layout(
             sharding=jax.sharding.NamedSharding(
                 jax_mesh, jax.sharding.PartitionSpec("batch", None)
             )
@@ -163,7 +172,7 @@ def test_distribute_variable_with_jax_layout(self):
         )
 
         variable = jax.numpy.array(np.random.normal(size=(16, 8)))
-        target_layout = jax_layout.Layout(
+        target_layout = self._create_jax_layout(
             sharding=jax.sharding.NamedSharding(
                 jax_mesh, jax.sharding.PartitionSpec("model", None)
             )
@@ -184,7 +193,7 @@ def test_distribute_input_data_with_jax_layout(self):
         )
 
         input_data = jax.numpy.array(np.random.normal(size=(16, 8)))
-        target_layout = jax_layout.Layout(
+        target_layout = self._create_jax_layout(
             sharding=jax.sharding.NamedSharding(
                 jax_mesh, jax.sharding.PartitionSpec("batch", None)
             )