jax

83d89b7a - [Pallas/Mosaic GPU] Introduce the `print_layout` primitive.

[Pallas/Mosaic GPU] Introduce the `print_layout` primitive. This new primitive allows printing the layout inferred by Mosaic GPU during lowering. PiperOrigin-RevId: 783186341