jax

7a007ea9 - [Pallas/Mosaic GPU] Add support for `s8` WGMMA with lhs in registers.

[Pallas/Mosaic GPU] Add support for `s8` WGMMA with lhs in registers. The low-level support only allows `swizzle=64` or larger for the time being. PiperOrigin-RevId: 826554680