jax

bf618a83 - [Mosaic:TPU] 16-bit tpu.iota

[Mosaic:TPU] 16-bit tpu.iota To support expressing iota over half-sublanes with a vreg shape of 8x128x2, for tpu.iota we now specify multiple dimensions. PiperOrigin-RevId: 775574142