jax

a046e21a - [pallas:mosaic_gpu] Do not do `mgpu.commit_shared` if all outputs are invariant wrt sequential axes