jax

147f6077 - [pallas:sc] Added additional checks for scatter/gather via `pltpu.async_copy`

[pallas:sc] Added additional checks for scatter/gather via `pltpu.async_copy` All of these are done at compilation time as well, but Python error messages are easier to understand than MLIR verification errors. PiperOrigin-RevId: 828447644