jax

a69deb89 - [Mosaic GPU][NFC] Rename `equations` to `eqns` in `gpu_layout_inference_test.py`.