jax

85193bdc - [NFC][Mosaic:TPU] Minor refactor of changeTiling to make chaining packs/unpacks and relayouts easier

[NFC][Mosaic:TPU] Minor refactor of changeTiling to make chaining packs/unpacks and relayouts easier Make the src layout mutable and get rid of src_offsets and src_tiling variables (get these directly from src) PiperOrigin-RevId: 775884959