{"record":{"id":"97a9e67d28d79bef","repo":"jax-ml/jax","slug":"derivatives-not-defined-for-partial-eigen-decompos","errorCode":null,"errorMessage":"Derivatives not defined for partial eigen decomposition.","messagePattern":"Derivatives not defined for partial eigen decomposition\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/linalg.py","lineNumber":1334,"sourceCode":"  avals_out = [v_aval, w_aval, info_aval]\n  rule = _linalg_ffi_lowering(target_name, avals_out=avals_out,\n                              operand_output_aliases={0: 0})\n  v, w, info = rule(ctx, operand, **kwargs)\n\n  zeros = mlir.full_like_aval(ctx, 0, info_aval)\n  ok = mlir.compare_hlo(info, zeros, \"EQ\", \"SIGNED\")\n  v = _replace_not_ok_with_nan(ctx, batch_dims, ok, v, v_aval)\n  w = _replace_not_ok_with_nan(ctx, batch_dims, ok, w, w_aval)\n  return [v, w]\n\n\ndef _eigh_jvp_rule(\n    primals, tangents, *, lower, sort_eigenvalues, subset_by_index, algorithm\n):\n  (a,) = primals\n  n = a.shape[-1]\n  if not (subset_by_index is None or subset_by_index == (0, n)):\n    raise NotImplementedError(\n        \"Derivatives not defined for partial eigen decomposition.\"\n    )\n  # Derivative for eigh in the simplest case of distinct eigenvalues.\n  # This is classic nondegenerate perurbation theory, but also see\n  # https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf\n  # The general solution treating the case of degenerate eigenvalues is\n  # considerably more complicated. Ambitious readers may refer to the general\n  # methods below or refer to degenerate perturbation theory in physics.\n  # https://www.win.tue.nl/analysis/reports/rana06-33.pdf and\n  # https://people.orie.cornell.edu/aslewis/publications/99-clarke.pdf\n  a_dot, = tangents\n\n  v, w_real = eigh_p.bind(\n      symmetrize(a),\n      lower=lower,\n      sort_eigenvalues=sort_eigenvalues,\n      subset_by_index=subset_by_index,\n      algorithm=algorithm,","sourceCodeStart":1316,"sourceCodeEnd":1352,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/linalg.py#L1316-L1352","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Differentiate the full-spectrum eigh and slice the result: w, v = eigh(a); use w[..., lo:hi] inside the loss","Use a nondifferentiable stop_gradient around the subset call if gradients through eigenpairs are not needed","Use an implicit/iterative method with known derivatives (LOBPCG-style custom VJP) for top-k training"],"exampleFix":"// before\nw, v = jax.lax.linalg.eigh(a, subset_by_index=(0, k))\nloss = f(w, v); jax.grad(loss)(a)  # raises\n// after\nw_all, v_all = jax.lax.linalg.eigh(a)\nw, v = w_all[..., :k], v_all[..., :, :k]\nloss = f(w, v); jax.grad(loss)(a)","handlingStrategy":"fallback","validationCode":"if subset_by_index not in (None,) and needs_grad:\n    subset_by_index = None  # differentiate full spectrum, slice in loss","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never combine subset_by_index with autodiff","Slice eigenpairs after full eigh instead"],"tags":["jax","autodiff","eigh","subset-by-index","gradients"],"backgroundTag":"nondifferentiable-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}