jax

8b02884c - [Mosaic:TPU] Fix trunc infer rule after cl/708011538

[Mosaic:TPU] Fix trunc infer rule after cl/708011538 For (1, 128) tiling 32-bit input, it assigns (1, 128) tiling at output, which can be invalid (e.g. it should be (1, 256) for bf16) PiperOrigin-RevId: 708112341