jax-ml/jax · error · NotImplementedError

Derivatives not defined for partial eigen decomposition.

Error message

Derivatives not defined for partial eigen decomposition.

What it means

jax/_src/lax/linalg.py:1334 in _eigh_jvp_rule. The JVP rule for eigh is only derived for the complete spectrum. If subset_by_index restricts the computed eigenpairs (anything other than None or (0, n)), differentiation raises NotImplementedError because the perturbation formula for partial decompositions is not implemented.

Source

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

  avals_out = [v_aval, w_aval, info_aval]
  rule = _linalg_ffi_lowering(target_name, avals_out=avals_out,
                              operand_output_aliases={0: 0})
  v, w, info = rule(ctx, operand, **kwargs)

  zeros = mlir.full_like_aval(ctx, 0, info_aval)
  ok = mlir.compare_hlo(info, zeros, "EQ", "SIGNED")
  v = _replace_not_ok_with_nan(ctx, batch_dims, ok, v, v_aval)
  w = _replace_not_ok_with_nan(ctx, batch_dims, ok, w, w_aval)
  return [v, w]


def _eigh_jvp_rule(
    primals, tangents, *, lower, sort_eigenvalues, subset_by_index, algorithm
):
  (a,) = primals
  n = a.shape[-1]
  if not (subset_by_index is None or subset_by_index == (0, n)):
    raise NotImplementedError(
        "Derivatives not defined for partial eigen decomposition."
    )
  # Derivative for eigh in the simplest case of distinct eigenvalues.
  # This is classic nondegenerate perurbation theory, but also see
  # https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf
  # The general solution treating the case of degenerate eigenvalues is
  # considerably more complicated. Ambitious readers may refer to the general
  # methods below or refer to degenerate perturbation theory in physics.
  # https://www.win.tue.nl/analysis/reports/rana06-33.pdf and
  # https://people.orie.cornell.edu/aslewis/publications/99-clarke.pdf
  a_dot, = tangents

  v, w_real = eigh_p.bind(
      symmetrize(a),
      lower=lower,
      sort_eigenvalues=sort_eigenvalues,
      subset_by_index=subset_by_index,
      algorithm=algorithm,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Differentiate the full-spectrum eigh and slice the result: w, v = eigh(a); use w[..., lo:hi] inside the loss
2. Use a nondifferentiable stop_gradient around the subset call if gradients through eigenpairs are not needed
3. Use an implicit/iterative method with known derivatives (LOBPCG-style custom VJP) for top-k training

Example fix

// before
w, v = jax.lax.linalg.eigh(a, subset_by_index=(0, k))
loss = f(w, v); jax.grad(loss)(a)  # raises
// after
w_all, v_all = jax.lax.linalg.eigh(a)
w, v = w_all[..., :k], v_all[..., :, :k]
loss = f(w, v); jax.grad(loss)(a)

if subset_by_index not in (None,) and needs_grad:
    subset_by_index = None  # differentiate full spectrum, slice in loss

When it happens

Trigger: Calling jax.grad / vjp on a function that uses jax.lax.linalg.eigh(..., subset_by_index=(lo, hi)) with a non-full range, i.e. combining partial-spectrum computation with autodiff. Only works at all on TPU (where subsets are supported), then fails at differentiation.

Common situations: End-to-end training of spectral layers (graph Laplacian eigenfilters, spectral clustering losses) where only k eigenpairs are computed to save memory, then backprop is attempted.

Related errors

• Derivatives of non-symmetric eigenvectors are only valid und
• subset_by_index not supported on CPU and GPU
• 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

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