jax-ml/jax · error · NotImplementedError

subset_by_index not supported on CPU and GPU

Error message

subset_by_index not supported on CPU and GPU

What it means

jax/_src/lax/linalg.py:1289 in _eigh_cpu_gpu_lowering. The subset_by_index argument of jax.lax.linalg.eigh (computing only eigenpairs [lo, hi)) is only implemented on TPU. On CPU and GPU the lowering only accepts the full range, so any other subset raises NotImplementedError.

Source

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

    )
  n = shape[0]
  d = (n if subset_by_index is None else
       subset_by_index[1] - subset_by_index[0])
  return (n, d), (d,)

def _eigh_dtype_rule(dtype, **_):
  return dtype, lax._complex_basetype(dtype)

def _eigh_cpu_gpu_lowering(
    ctx, operand, *, lower, sort_eigenvalues, subset_by_index, algorithm,
    target_name_prefix: str
):
  del sort_eigenvalues  # The CPU/GPU implementations always sort.
  operand_aval, = ctx.avals_in
  v_aval, w_aval = ctx.avals_out
  n = operand_aval.shape[-1]
  if not (subset_by_index is None or subset_by_index == (0, n)):
    raise NotImplementedError("subset_by_index not supported on CPU and GPU")
  batch_dims = operand_aval.shape[:-2]

  if algorithm == EighImplementation.QDWH:
    raise NotImplementedError("QDWH implementation is only supported on TPU")
  if algorithm == EighImplementation.JACOBI and target_name_prefix == "cpu":
    raise NotImplementedError("Jacobi implementation is not supported on CPU")

  if target_name_prefix == "cpu":
    dtype = operand_aval.dtype
    prefix = "he" if dtypes.issubdtype(dtype, np.complexfloating) else "sy"
    target_name = lapack.prepare_lapack_call(f"{prefix}evd_ffi",
                                             operand_aval.dtype)
    kwargs = {
      "mode": np.uint8(ord("V")),
      "uplo": np.uint8(ord("L" if lower else "U")),
    }
  else:
    target_name = f"{target_name_prefix}solver_syevd_ffi"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Compute the full spectrum and slice: evals, evecs = eigh(a); take evals[lo:hi] (fine when n is moderate)
  2. For top-k on GPU/CPU use an iterative method (e.g. jax.scipy.sparse.linalg.lobpcg) or scipy.sparse.linalg.eigsh on the host
  3. Run on TPU backend where subset_by_index is supported

Example fix

// before
evals, evecs = jax.lax.linalg.eigh(a, subset_by_index=(0, 10))  # on GPU
// after
evals_all, evecs_all = jax.lax.linalg.eigh(a)
evals, evecs = evals_all[..., :10], evecs_all[..., :, :10]
Defensive patterns

Strategy: fallback

Validate before calling

backend = jax.default_backend()
use_subset = backend == 'tpu' and lo == 0
if not use_subset:
    plan = 'full eigh + slice'

Try / catch

try:
    w, v = jax.lax.linalg.eigh(a, subset_by_index=(lo, hi))
except NotImplementedError:
    w_all, v_all = jax.lax.linalg.eigh(a)
    w, v = w_all[..., lo:hi], v_all[..., :, lo:hi]

Prevention

When it happens

Trigger: Calling jax.lax.linalg.eigh(a, subset_by_index=(lo, hi)) with a range other than (0, n) or None while running on CPU or GPU backend. Note lower bounds other than 0 are also unsupported even on TPU.

Common situations: Porting scipy.linalg.eigh(..., subset_by_index=...) partial-spectrum code to JAX; memory-saving attempts to compute only the top-k eigenvalues of large matrices; code that ran on TPU moved to a GPU box.

Related errors

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