jax

7db3a70e - [Pallas:MGPU] Add support for batch dimensions to `async_store_tmem`.

[Pallas:MGPU] Add support for batch dimensions to `async_store_tmem`. When batch dimensions exist, we iterate over all indices in the batch shape, extract the corresponding slice from the vector registers, compute the correct column offset in TMEM and perform sliced subview stores. PiperOrigin-RevId: 932455324