jax

caa0b9ef - [mosaic_gpu] Support TMA_INDICES_4_LAYOUT for TMA gather/scatter (replicated).

[mosaic_gpu] Support TMA_INDICES_4_LAYOUT for TMA gather/scatter (replicated). Adds a new layout `TMA_INDICES_4_LAYOUT` where every thread holds the same indices. Updates `launch_context.py` to support this layout in TMA gather and scatter operations by adjusting the thread predicate to `WARPGROUP` instead of `WARP`, and fixing the row index stride. Parameterized existing gather and scatter tests to cover both layouts. PiperOrigin-RevId: 925455297