jax

f14c260e - [Mosaic GPU][NFC] Import `equations` as `eqns` in `layout_inference2.py`.