I use VJP frequently in my project. It runs the function that is subject to Jacobian computation and returns a primals_out together with the callable vjp function. For example, custom VJP definition in JAX documentation is given like this:
from jax import custom_vjp
@custom_vjp
def f(x, y):
return jnp.sin(x) * y
def f_fwd(x, y):
# Returns primal output and residuals to be used in backward pass by f_bwd.
return f(x, y), (jnp.cos(x), jnp.sin(x), y)
def f_bwd(res, g):
cos_x, sin_x, y = res # Gets residuals computed in f_fwd
return (cos_x * g * y, sin_x * g)
f.defvjp(f_fwd, f_bwd)
In this example, we see that evaluation of the forward function is required when using VJP. This is also the case when using regular VJP instead of a custom defined one. However, when the evaluation of function costs highly and since I have already run that function somewhere in my code, I don't want VJP to evaluate that function one more time.
So, is there any way to indicate that a function will not be evaluated when computing its VJP?