Add "format" support for latest JAX. (#21406)

The `jax.experimental.layout.Format` configuration has a `"format"` attribute
rather than a `"layout"` attribute.

Author: cantonios

keras/src/backend/jax/distribution_lib.py

@@ -78,11 +78,16 @@ def distribute_tensor(tensor, layout):
             layout, jax.sharding.Sharding
         ) and tensor.sharding.is_equivalent_to(layout, ndim=len(tensor.shape)):
             return tensor
-        # JAX explicit layout support.
+        # JAX explicit "layout" support.
         elif hasattr(layout, "layout"):
             current_layout = getattr(tensor, "layout", None)
             if current_layout == layout:
                 return tensor
+        # JAX explicit "format" support.
+        elif hasattr(layout, "format"):
+            current_layout = getattr(tensor, "format", None)
+            if current_layout == layout:
+                return tensor
 
     return jax.device_put(tensor, layout)