jax

ed952c8e - [Pallas TPU] Support jnp.take_along_axis for 32-bit vreg-sized vector.