jax

78520455 - [Mosaic TPU] Enable non-sublane-aligned bf16 2D load/stores for earlier TPU gens

[Mosaic TPU] Enable non-sublane-aligned bf16 2D load/stores for earlier TPU gens It is still not efficiently implemented, this is mostly to clean up some logic. We may be able to fuse the creation of masks for different tiles into the creation of a single one. But this is also a problem for the later gens. This also cleans up an unreachable return statement. PiperOrigin-RevId: 714847066