Update explicit layouts to use JAX 0.6.2 layout.Format. (#125)

There was a rename in JAX from `jax.experimental.layout.Layout`
to `jax.experimental.layout.Format`, and JAX arrays have an
attribute change `jax.Array.layout -> jax.Array.format`.

Author: cantonios

keras_rs/src/layers/embedding/distributed_embedding_test.py

@@ -560,10 +560,10 @@ def test_correctness(
                 # Determine explicit shardings/layouts for jit compilation
                 # (required for sparsecore computations).
                 trainable_layouts = keras.tree.map_structure(
-                    lambda x: x.value.layout, layer.trainable_variables
+                    lambda x: x.value.format, layer.trainable_variables
                 )
                 non_trainable_layouts = keras.tree.map_structure(
-                    lambda x: x.value.layout, layer.non_trainable_variables
+                    lambda x: x.value.format, layer.non_trainable_variables
                 )
                 # Input/output data involved in sparsecore operations are
                 # sharded across all sparse-core-capable devices.

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.2
+jax[cuda12]==0.6.2
 
 # Support for large embeddings.
 jax-tpu-embedding;sys_platform == 'linux' and platform_machine == 'x86_64'

requirements.txt

@@ -7,7 +7,7 @@ tensorflow~=2.18.0;sys_platform == 'darwin'
 torch>=2.1.0
 
 # Jax.
-jax[cpu]
+jax[cpu]>=0.6.2
 jax-tpu-embedding;sys_platform == 'linux' and platform_machine == 'x86_64'
 
 # pre-commit checks (formatting, linting, etc.)