Why Error: "TypeError: can't apply forward-mode autodiff (jvp) to a custom_vjp function."

[miaoL1]

Learned from @yonghakim, I utilized code below:

def conj(arr):
    return arr.real + arr.imag * -1j

@partial(jax.custom_vjp, nondiff_argnums=(1, 2, 3))
def eig(x, type_complex=jnp.complex128, perturbation=1E-10, device='cpu'):
    _eig = jax.jit(jnp.linalg.eigh, device=jax.devices('cpu')[0])
    eigenvalues_shape = jax.ShapeDtypeStruct(x.shape[:-1], type_complex)
    eigenvectors_shape = jax.ShapeDtypeStruct(x.shape, type_complex)
    result_shape_dtype = (eigenvalues_shape, eigenvectors_shape)
    if device == 'cpu':
        res = _eig(x)
    else:
        res = jax.pure_callback(_eig, result_shape_dtype, x)
    return res

def eig_fwd(x, type_complex, perturbation, device):
    return eig(x, type_complex, perturbation), (eig(x, type_complex, perturbation), x)

def eig_bwd(type_complex, perturbation, device, res, g):
    (eig_val, eig_vector), x = res
    grad_eigval, grad_eigvec = g
    grad_eigval = jnp.diag(grad_eigval)
    W_H = eig_vector.T.conj()
    Fij = eig_val.reshape((1, -1)) - eig_val.reshape((-1, 1))
    Fij = Fij / (jnp.abs(Fij) ** 2 + perturbation)
    Fij = Fij.at[jnp.diag_indices_from(Fij)].set(0)
    grad = jnp.linalg.inv(W_H) @ (grad_eigval.conj() + Fij * (W_H @ grad_eigvec.conj())) @ W_H
    grad = grad.conj()
    if not jnp.iscomplexobj(x):
        grad = grad.real
    return grad,

eig.defvjp(eig_fwd, eig_bwd)

to conduct gradient calculation on degenerated eigenvectors (which I've cried for help for few days ago #19030) .

Though the gradient is accessible, when I tried to optimize parameters, error occured:

File D: ...\lib\site-packages\jax\_src\interpreters\ad.py:751, in raise_custom_vjp_error_on_jvp(*_, **__)
    750 def raise_custom_vjp_error_on_jvp(*_, **__):
--> 751   raise TypeError("can't apply forward-mode autodiff (jvp) to a custom_vjp "
    752                   "function.")

TypeError: can't apply forward-mode autodiff (jvp) to a custom_vjp function.

I'm new to JAX and my question might be ambiguous or inappropriate, but still I would really appreciate it if anyone can offer me cludes.

[jakevdp]

Hi - thanks for the question! In JAX, forward-mode autodiff is handled by jvp (Jacobian vector product), while reverse-mode autodiff is handled by vjp (vector Jacobian product).

When you define custom_vjp, you are defining how reverse-mode autodiff should behave, but you have not defined how forward-mode autodiff should behave.

To fix this, you can either stick to reverse-mode autodiff transformations (i.e. use vjp, jacrev, jacobian, or grad, all of which use reverse-mode autodiff) or you can define a custom_jvp for your function, which defines how it should behave with forward-mode autodiff.

For more details, see Custom derivative rules for Python code.

[miaoL1]

Thanks for enlightening me!