jax-ml/jax · error · NotImplementedError

Derivatives of non-symmetric eigenvectors are only valid und

Error message

Derivatives of non-symmetric eigenvectors are only valid under assumptions on the input that JAX cannot check (see the enable_eigvec_derivs argument to jax.lax.linalg.eig). Pass enable_eigvec_derivs=True to jax.lax.linalg.eig to opt in. See https://github.com/jax-ml/jax/issues/2748 for discussion.

What it means

jax/_src/lax/linalg.py:1228 in eig_jvp_rule. JAX refuses to differentiate jax.lax.linalg.eig through eigenvectors by default: the derivative formula v' = (A - λI)^+ ... is ill-conditioned/undefined for degenerate (repeated) eigenvalues, and JAX cannot statically verify the required assumptions. You must explicitly opt in with enable_eigvec_derivs=True.

Source

Thrown at jax/_src/lax/linalg.py:1228

  U = dot(v, Fmat * P)
  # The eigenvalue equation gives dv_j = u_j + c_j v_j with c_j free; the two
  # real LAPACK normalisation constraints fix c_j to
  #   c_j = -Re(v_j* . u_j) - i Im(u_{k_j j}) / v_{k_j j},  k_j = argmax_i |v_ij|.
  k = lax.argmax(lax.abs(v), axis=v.ndim - 2, index_dtype=np.int32)
  mask = (lax.broadcasted_iota(np.int32, v.shape, v.ndim - 2)
          == lax.expand_dims(k, (v.ndim - 2,))).astype(v.dtype)
  c = lax.complex(-(v.conj() * U).sum(-2).real,
                  -(mask * U).sum(-2).imag / (mask * v).sum(-2).real)
  return dw, U + v * lax.expand_dims(c, (v.ndim - 2,))

def eig_jvp_rule(primals, tangents, *, compute_left_eigenvectors,
                 compute_right_eigenvectors, enable_eigvec_derivs,
                 implementation):
  a, = primals
  da, = tangents
  if compute_left_eigenvectors or compute_right_eigenvectors:
    if not enable_eigvec_derivs:
      raise NotImplementedError(
          'Derivatives of non-symmetric eigenvectors are only valid under '
          'assumptions on the input that JAX cannot check (see the '
          'enable_eigvec_derivs argument to jax.lax.linalg.eig). Pass '
          'enable_eigvec_derivs=True to jax.lax.linalg.eig to opt in. See '
          'https://github.com/jax-ml/jax/issues/2748 for discussion.')
  outs = eig(a, compute_left_eigenvectors=compute_left_eigenvectors,
             compute_right_eigenvectors=True,
             enable_eigvec_derivs=enable_eigvec_derivs,
             implementation=implementation)
  w, vr = outs[0], outs[-1]
  dot = partial(lax.dot if a.ndim == 2 else lax.batch_matmul,
                precision=lax.Precision.HIGHEST)
  da = da.astype(vr.dtype)
  if not (compute_left_eigenvectors or compute_right_eigenvectors):
    return [w], [(_solve(vr, da) * _T(vr)).sum(-1)]
  dw, dvr = _eig_vec_jvp(dot, w, vr, da)
  primal_out, tangent_out = [w], [dw]
  if compute_left_eigenvectors:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Opt in if your inputs are known to have distinct eigenvalues: jax.lax.linalg.eig(a, enable_eigvec_derivs=True)
2. Differentiate only eigenvalues (compute_*_eigenvectors=False) which have well-defined derivatives
3. Use eigh instead if the matrix is symmetric/Hermitian — its eigenvector derivatives are the standard well-defined case
4. For robustness, add a small regularization/penalty to keep the spectrum separated, or use a differentiable surrogate (e.g. power iteration / LOBPCG-style layers)

Example fix

// before
w, v = jax.lax.linalg.eig(a)
loss = f(v)
jnp.gradient... jax.grad(loss)(a)  # raises
// after
w, v = jax.lax.linalg.eig(a, enable_eigvec_derivs=True)  # only if eigenvalues are distinct

if grads_needed and eigenvectors_used:
    w, v = jax.lax.linalg.eig(a, enable_eigvec_derivs=True)  # only if spectrum distinct
else:
    w, v = jax.lax.linalg.eig(a, compute_left_eigenvectors=False)

try:
    jax.grad(f)(a)
except NotImplementedError:
    f = jax.tree_util.Partial(f, enable_eigvec_derivs=True)
    jax.grad(f)(a)

When it happens

Trigger: Calling jax.grad / jax.jit(jax.vjp(...)) on a function whose output includes eigenvectors of jax.lax.linalg.eig (or jnp.linalg.eig) without enable_eigvec_derivs=True; happens with compute_left_eigenvectors or compute_right_eigenvectors True (right is default).

Common situations: Differentiating through spectral decompositions in physics-informed ML, graph networks using eigenvectors of a learned matrix, PCA-like layers trained end-to-end. Silent math validity issue: near-degenerate spectra produce wildly wrong gradients.

Related errors

• Derivatives not defined for partial eigen decomposition.
• primal and tangent arguments to jax.jvp must be tuples or li
• primal and tangent arguments to jax.jvp must have the same t
• primal and tangent arguments to jax.jvp do not match; dtypes
• jvp called with different primal and tangent shapes;Got prim

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/ae39abe887c338b0. Report an issue: GitHub.