jax

823abec3 - Second attempt at fixing Mosaic GPU CI failures after bzlmod switch.