jax

dfb2a19c - [Mosaic TPU] Support reshapes of form `[m, n] -> [m*n, 1]` and `[m*n] -> [m, n, 1]` where `n % 128 == 0`.