jax

4a596aee - [Mosaic TPU] Force offset to 0 when inferring input has offset out of the first tile.

[Mosaic TPU] Force offset to 0 when inferring input has offset out of the first tile. We still have this temporary check in apply vector layout, but in infer vector layout, instead of throwing error, we should just reset offset to zero. Because some ops which has relaxed this restriction might be passed as input for un-relaxed ops and cause failure. PiperOrigin-RevId: 680706301