jax-ml/jax · error · NotImplementedError

Only the eigvals=None case of eigh is implemented.

Error message

Only the eigvals=None case of eigh is implemented.

What it means

jax.scipy.linalg.eigh does not support the eigvals subset parameter; only eigvals=None (all eigenvalues) is implemented. SciPy uses eigvals=(lo, hi) to return a range of eigenvalues/eigenvectors, which requires dynamic-shape support JAX lacks.

Source

Thrown at jax/_src/scipy/linalg.py:403

          eigvals: None, type: int) -> Array: ...

@overload
def _eigh(a: ArrayLike, b: ArrayLike | None, lower: bool, eigvals_only: Literal[False],
          eigvals: None, type: int) -> tuple[Array, Array]: ...

@overload
def _eigh(a: ArrayLike, b: ArrayLike | None, lower: bool, eigvals_only: bool,
          eigvals: None, type: int) -> Array | tuple[Array, Array]: ...

@jit(static_argnames=('lower', 'eigvals_only', 'eigvals', 'type'))
def _eigh(a: ArrayLike, b: ArrayLike | None, lower: bool, eigvals_only: bool,
          eigvals: None, type: int) -> Array | tuple[Array, Array]:
  if b is not None:
    raise NotImplementedError("Only the b=None case of eigh is implemented")
  if type != 1:
    raise NotImplementedError("Only the type=1 case of eigh is implemented.")
  if eigvals is not None:
    raise NotImplementedError(
        "Only the eigvals=None case of eigh is implemented.")

  a, = promote_dtypes_inexact(jnp.asarray(a))
  v, w = lax_linalg.eigh(a, lower=lower)

  if eigvals_only:
    return w
  else:
    return w, v

@overload
def eigh(a: ArrayLike, b: ArrayLike | None = None, lower: bool = True,
         eigvals_only: Literal[False] = False, overwrite_a: bool = False,
         overwrite_b: bool = False, turbo: bool = True, eigvals: None = None,
         type: int = 1, check_finite: bool = True) -> tuple[Array, Array]: ...

@overload
def eigh(a: ArrayLike, b: ArrayLike | None = None, lower: bool = True, *,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Compute all eigenpairs, then slice: w, v = eigh(a); w = w[:k]; v = v[:, :k]
2. Use an iterative method better suited to partial spectra, e.g. JAX implementations of Lanczos/LOBPCG (jax.scipy.sparse.linalg.lobpcg)
3. Call scipy.linalg.eigh host-side outside jit if GPU execution is unnecessary

Example fix

// before
w, v = jax.scipy.linalg.eigh(a, eigvals=(0, 2))
// after
w, v = jax.scipy.linalg.eigh(a)
w, v = w[:3], v[:, :3]

assert eigvals is None, 'jax eigh does not support eigvals; slice results instead'

null

When it happens

Trigger: Calling jax.scipy.linalg.eigh(a, eigvals=(0, 2)) or any non-None eigvals tuple.

Common situations: Porting SciPy spectral code (e.g. computing only lowest-k eigenpairs for dimensionality reduction or graph Laplacians) to JAX and passing the SciPy-style eigvals argument.

Related errors

• Only the type=1 case of eigh is implemented.
• eigh_tridiagonal(..., select='v') is not implemented
• Blocked version is not implemented yet.
• array ref with memory space only works inside of a `jit`.
• pinned array ref only works inside of a `jit`.

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