has_aux custom.jvp

[pharringtonp19]

Is it possible to pass in additional arguments to custom.jvp? Edit: I believe this would be related to 497 and 499

In my particular case, I don't want to not compute primal_out but instead would like to pass in a value that would take its place. Something like the following:

import jax.numpy as jnp
from jax import custom_jvp

@custom_jvp
def f(x, y):
  return jnp.sin(x) * y

@f.defjvp
def f_jvp(primals, tangents, psuedo_primals_out):
  x, y = primals
  x_dot, y_dot = tangents
  #primal_out = f(x, y)  Don't compute primal_out!
  tangent_out = jnp.cos(x) * x_dot * y + jnp.sin(x) * y_dot
  return psuedo_primals_out, tangent_out