flax

dcd2f5e2 - Use `.input_formats` and `.output_formats` in place of `.input_layouts` and `.output_layouts` respectively.

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