jax

3366c927 - [Mosaic GPU][NFC] Simplify and clean up layout inference tests to use `FuncOp`s.