jax

e59303cf - [Mosaic GPU] Simplify the matmul example

[Mosaic GPU] Simplify the matmul example Remove a bunch of WGMMAImpl classes. This is meant to be a simple forkable example, not a complete kernel. PiperOrigin-RevId: 655923069