jax

7a46165b - [pallas:mosaic] Yet another fix for named grid and batching