Is there a way to disable forward evaluation while using VJP?

[aonurdasdemir]

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?

[shoyer]

JAX needs to evaluate the forward pass in order to calculate intermediate quantitites that are used in the VJP calculation. So there's not really any avoiding it.

Probably your best option is to use jax.vjp (or jax.value_and_grad) for both your evaluation of the forward pass and the backwards pass. Or if your function happens to be linear (in which case intermediate values from the forward pass aren't needed), you could put the call to vjp under a jit and simply ignore the forward pass result, in which case XLA should remove the unused computation.

[llCurious]

hey, i have similar problem. I try to use jit and i found that the unused computations are not removed from the compiled XLA.

BTW, my purpose is make use of the derivates of primitive functions and i need to use vjp to pass the gradients.

def custom_Dot_bwd(res, g):
  print("Custom Dot backward", "="*20)
  inputs, W = res
  # custom modification to inputs and W
  # ...

  # manually compute the gradients of inputs and W
  # return jnp.dot(g, W_hi.T), jnp.dot(inputs_hi.T, g)

  # automatically compute the gradients, pass g to the vjp. However, the forward pass seems to be inevitable
  _, grad_fn = jax.vjp(jnp.dot, inputs_hi, W_hi)
  return grad_fn(g)

If i try to custom convolution, the first method seems to be not general. While the second method requires to compute the forward again. Do you have any suggestions?