jax

ef4dbd9c - [Mosaic TPU] Add support for packing to 16-bit integers on TPUv4

[Mosaic TPU] Add support for packing to 16-bit integers on TPUv4 And refactor some test conditions to better match what we really support. The tests were failing on older TPUs. PiperOrigin-RevId: 716214098