add handlers for jax.custom_jvp / jax.custom_vjp from github.com/google/jax/pull/2500
add handlers for jax.custom_jvp / jax.custom_vjp from github.com/google/jax/pull/2500

PiperOrigin-RevId: 303682287
Change-Id: Ibd516de01647b8a1aa2518a09d47510024ac7511

deepmind/dm-haiku,deepmind/dm-haiku