jax

66476a0b - [Mosaic GPU] Add layout inference and lowering for `scf.WhileOp` and enable tests.