Use Layout in place of DeviceLocalLayout and .layout in place of .device_local_layout.

JAX is undergoing a rename of the contents of jax.experimental.layouts in preparation for its graduation from experimental:

* "Formats" are replacing "layouts", and specifically `Layout -> Format`
* "Layouts" are replacing "device local layouts", and specifically `DeviceLocalLayout -> Layout`

This is an incremental update carrying out #2.

PiperOrigin-RevId: 774920904

Author: yashk2810

recml/layers/linen/sparsecore.py

@@ -30,6 +30,11 @@
 from recml.core.ops import embedding_ops
 import tensorflow as tf
 
+if jax.__version_info__ >= (0, 6, 3):
+  DLL = layout.Layout
+else:
+  DLL = layout.DeviceLocalLayout  # type: ignore
+
 
 with epy.lazy_imports():
   # pylint: disable=g-import-not-at-top
@@ -382,7 +387,7 @@ class SparsecoreLayout(nn.Partitioned[A]):
   def get_sharding(self, _):
     assert self.mesh is not None
     return layout.Format(
-        layout.DeviceLocalLayout(major_to_minor=(0, 1), _tiling=((8,),)),
+        DLL(major_to_minor=(0, 1), _tiling=((8,),)),
         jax.sharding.NamedSharding(self.mesh, self.get_partition_spec()),
     )