jax

0d2f49ab - [Mosaic GPU] Generalize f8 casting routines

[Mosaic GPU] Generalize f8 casting routines In particular, add support for casting to and from ue8m0 and f32/bf16. Also, add support for casts for layouts that don't have an even vector length. PiperOrigin-RevId: 859569741