jax

0b6c3550 - [Pallas] Add experimental (private for now) API for manual fusion into Pallas kernels