jax

d2bf034c - [Mosaic GPU] Test the wgmma_op lowering when a is in registers.

[Mosaic GPU] Test the wgmma_op lowering when a is in registers. I had to add support for wgmma layout in vector_load. Not sure if this is useful outside the test. PiperOrigin-RevId: 735384104