jax

74938be8 - Consolidate initial/final style custom_vjp primitives into one.