jax

bdc7cf04 - [Mosaic GPU] Enable bidirectional transpose relayouts in `FragmentedArray.to_layout`.

[Mosaic GPU] Enable bidirectional transpose relayouts in `FragmentedArray.to_layout`. This change allows `FragmentedArray.to_layout` to convert between `WGMMA_LAYOUT` and `WGMMA_TRANSPOSED_LAYOUT`, and between `TCGEN05_LAYOUT` and `TCGEN05_TRANSPOSED_LAYOUT` in both directions. Previously, only one direction was supported. PiperOrigin-RevId: 829394697