jax

d48009f6 - [Mosaic] Use small second minor tiling for packed type reshape of contiguous vreg slice.

[Mosaic] Use small second minor tiling for packed type reshape of contiguous vreg slice. We can turn reshape into a row shuffle inside a single vreg with small second minor tiling if **there is no padding and each vreg holds a contiguous slice of the flattened data**. The strategy is to choose tiling = `(elements_per_vreg // shape[-1], target_shape[1])` and make sure each tile is full without padding. For example, bf16[16, 128] with tiling [16, 128] -> bf16[8, 256] with tiling [8, 128]. This can be done by mixed unpacking/packing to halve/double sublane tiling. Alternatively, we can unpack to 32 bit and use sublane shuffle, but we lack of support for general tpu::SublaneShuffleOp. And it's not going to be faster IMO as arbitrary shuffle pattern needs to be emulated by store + load. PiperOrigin-RevId: 817620266