Use .format in place of .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 #1.

PiperOrigin-RevId: 773116112

Author: yashk2810

flax/nnx/transforms/compilation.py

@@ -599,7 +599,7 @@ def output_shardings(self):  # PyTree[sharding.Sharding]
 
   @property
   def input_layouts(self):
-    return self.compiled.input_layouts
+    return self.compiled.input_formats
 
 
 @dataclasses.dataclass(frozen=True, slots=True)