jax-ml/jax · error · NotImplementedError

Only the type=1 case of eigh is implemented.

Error message

Only the type=1 case of eigh is implemented.

What it means

jax.scipy.linalg.eigh only implements the standard eigenproblem (type=1). The `type` parameter selects which generalized eigenproblem formulation to solve per LAPACK's syevd/heevd family; JAX has only implemented type 1, so any other integer raises NotImplementedError.

Source

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

@overload
def _eigh(a: ArrayLike, b: ArrayLike | None, lower: bool, eigvals_only: Literal[True],
          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]: ...

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Drop the type argument (use default type=1) if you only need the ordinary symmetric eigenproblem
  2. For generalized problems, reduce manually: solve eigh(solve(cholesky(b), a)) via a congruence transform with b's Cholesky factor, then map eigenvectors back
  3. Use scipy.linalg.eigh on the host (outside jit) if GPU execution is not required
  4. File/track a feature request on the JAX GitHub repo

Example fix

// before
w, v = jax.scipy.linalg.eigh(a, b, type=2)
// after
L = jax.scipy.linalg.cholesky(b)
a_whitened = jax.scipy.linalg.solve_triangular(L, a, lower=True)
w, v = jax.scipy.linalg.eigh(a_whitened.T + a_whitened)  # symmetrize as needed
v = jax.scipy.linalg.solve_triangular(L.T, v)
Defensive patterns

Strategy: validation

Validate before calling

if type != 1: raise ValueError('jax eigh supports only type=1')

Type guard

null

Prevention

When it happens

Trigger: Calling jax.scipy.linalg.eigh(a, type=2) or type=3 (also requires b != None, which itself raises a sibling error).

Common situations: Porting NumPy/SciPy code that uses scipy.linalg.eigh(a, b, type=2) for generalized eigenvalue problems (e.g. vibronic analysis, LDA/QDA discriminants) to JAX.

Related errors


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