{"record":{"id":"ae39abe887c338b0","repo":"jax-ml/jax","slug":"derivatives-of-non-symmetric-eigenvectors-are-only","errorCode":null,"errorMessage":"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.","messagePattern":"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\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/linalg.py","lineNumber":1228,"sourceCode":"  U = dot(v, Fmat * P)\n  # The eigenvalue equation gives dv_j = u_j + c_j v_j with c_j free; the two\n  # real LAPACK normalisation constraints fix c_j to\n  #   c_j = -Re(v_j* . u_j) - i Im(u_{k_j j}) / v_{k_j j},  k_j = argmax_i |v_ij|.\n  k = lax.argmax(lax.abs(v), axis=v.ndim - 2, index_dtype=np.int32)\n  mask = (lax.broadcasted_iota(np.int32, v.shape, v.ndim - 2)\n          == lax.expand_dims(k, (v.ndim - 2,))).astype(v.dtype)\n  c = lax.complex(-(v.conj() * U).sum(-2).real,\n                  -(mask * U).sum(-2).imag / (mask * v).sum(-2).real)\n  return dw, U + v * lax.expand_dims(c, (v.ndim - 2,))\n\ndef eig_jvp_rule(primals, tangents, *, compute_left_eigenvectors,\n                 compute_right_eigenvectors, enable_eigvec_derivs,\n                 implementation):\n  a, = primals\n  da, = tangents\n  if compute_left_eigenvectors or compute_right_eigenvectors:\n    if not enable_eigvec_derivs:\n      raise NotImplementedError(\n          'Derivatives of non-symmetric eigenvectors are only valid under '\n          'assumptions on the input that JAX cannot check (see the '\n          'enable_eigvec_derivs argument to jax.lax.linalg.eig). Pass '\n          'enable_eigvec_derivs=True to jax.lax.linalg.eig to opt in. See '\n          'https://github.com/jax-ml/jax/issues/2748 for discussion.')\n  outs = eig(a, compute_left_eigenvectors=compute_left_eigenvectors,\n             compute_right_eigenvectors=True,\n             enable_eigvec_derivs=enable_eigvec_derivs,\n             implementation=implementation)\n  w, vr = outs[0], outs[-1]\n  dot = partial(lax.dot if a.ndim == 2 else lax.batch_matmul,\n                precision=lax.Precision.HIGHEST)\n  da = da.astype(vr.dtype)\n  if not (compute_left_eigenvectors or compute_right_eigenvectors):\n    return [w], [(_solve(vr, da) * _T(vr)).sum(-1)]\n  dw, dvr = _eig_vec_jvp(dot, w, vr, da)\n  primal_out, tangent_out = [w], [dw]\n  if compute_left_eigenvectors:","sourceCodeStart":1210,"sourceCodeEnd":1246,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/linalg.py#L1210-L1246","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Opt in if your inputs are known to have distinct eigenvalues: jax.lax.linalg.eig(a, enable_eigvec_derivs=True)","Differentiate only eigenvalues (compute_*_eigenvectors=False) which have well-defined derivatives","Use eigh instead if the matrix is symmetric/Hermitian — its eigenvector derivatives are the standard well-defined case","For robustness, add a small regularization/penalty to keep the spectrum separated, or use a differentiable surrogate (e.g. power iteration / LOBPCG-style layers)"],"exampleFix":"// before\nw, v = jax.lax.linalg.eig(a)\nloss = f(v)\njnp.gradient... jax.grad(loss)(a)  # raises\n// after\nw, v = jax.lax.linalg.eig(a, enable_eigvec_derivs=True)  # only if eigenvalues are distinct","handlingStrategy":"validation","validationCode":"if grads_needed and eigenvectors_used:\n    w, v = jax.lax.linalg.eig(a, enable_eigvec_derivs=True)  # only if spectrum distinct\nelse:\n    w, v = jax.lax.linalg.eig(a, compute_left_eigenvectors=False)","typeGuard":null,"tryCatchPattern":"try:\n    jax.grad(f)(a)\nexcept NotImplementedError:\n    f = jax.tree_util.Partial(f, enable_eigvec_derivs=True)\n    jax.grad(f)(a)","preventionTips":["Check eigenvalue separation before training through eigenvectors","Prefer eigh for symmetric matrices","Consider stop_gradient around eigenvectors when only used as features"],"tags":["jax","autodiff","eig","gradients","numerical"],"backgroundTag":"nondifferentiable-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}