jax

00d9f452 - [Pallas/Fuser] Add support for custom_call_jvp/pjit to push_block_spec