Use .input_formats and .output_formats in place of .input_layouts and .output_layouts respectively.

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: 773109310

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)