Update references of jax_layout.Layout to jax_layout.Format. (#124)

This reflects an upstream change in JAX.

Author: cantonios

keras_rs/src/layers/embedding/jax/distributed_embedding.py

@@ -36,7 +36,7 @@
 def _get_partition_spec(
     layout: (
         keras.distribution.TensorLayout
-        | jax_layout.Layout
+        | jax_layout.Format
         | jax.sharding.NamedSharding
         | jax.sharding.PartitionSpec
     ),
@@ -45,7 +45,7 @@ def _get_partition_spec(
     if isinstance(layout, keras.distribution.TensorLayout):
         layout = layout.backend_layout
 
-    if isinstance(layout, jax_layout.Layout):
+    if isinstance(layout, jax_layout.Format):
         layout = layout.sharding
 
     if isinstance(layout, jax.sharding.NamedSharding):
@@ -217,7 +217,7 @@ def _create_sparsecore_distribution(
         sparsecore_layout = keras.distribution.TensorLayout(axes, device_mesh)
         # Custom sparsecore layout with tiling.
         # pylint: disable-next=protected-access
-        sparsecore_layout._backend_layout = jax_layout.Layout(
+        sparsecore_layout._backend_layout = jax_layout.Format(
             jax_layout.DeviceLocalLayout(
                 major_to_minor=(0, 1),
                 _tiling=((8,),),

keras_rs/src/layers/embedding/jax/distributed_embedding_test.py

@@ -41,7 +41,7 @@ def _create_sparsecore_layout(
     )
     sparsecore_layout = keras.distribution.TensorLayout(axes, device_mesh)
     # Custom sparsecore layout with tiling.
-    sparsecore_layout._backend_layout = jax_layout.Layout(  # pylint: disable=protected-access
+    sparsecore_layout._backend_layout = jax_layout.Format(  # pylint: disable=protected-access
         jax_layout.DeviceLocalLayout(
             major_to_minor=(0, 1),
             _tiling=((8,),),

keras_rs/src/layers/embedding/jax/embedding_lookup.py

@@ -8,7 +8,7 @@
 
 import jax
 import numpy as np
-from jax.experimental import layout
+from jax.experimental import layout as jax_layout
 from jax_tpu_embedding.sparsecore.lib.nn import embedding
 from jax_tpu_embedding.sparsecore.lib.nn import embedding_spec
 from jax_tpu_embedding.sparsecore.utils import utils as jte_utils
@@ -20,7 +20,7 @@
 shard_map = jax.experimental.shard_map.shard_map  # type: ignore[attr-defined]
 
 ArrayLike: TypeAlias = jax.Array | np.ndarray[Any, Any]
-JaxLayout: TypeAlias = jax.sharding.NamedSharding | layout.Layout
+JaxLayout: TypeAlias = jax.sharding.NamedSharding | jax_layout.Format
 
 
 class EmbeddingLookupConfiguration:

requirements-jax-cuda.txt

@@ -8,7 +8,7 @@ torch>=2.1.0
 # Jax with cuda support.
 # Keep same version as Keras repo.
 --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-jax[cuda12_pip]==0.6.0
+jax[cuda12_pip]==0.6.2
 
 # Support for large embeddings.
 jax-tpu-embedding;sys_platform == 'linux' and platform_machine == 'x86_64'