jax-ml/jax · error · NotImplementedError

eigh_tridiagonal(..., select='v') is not implemented

Error message

eigh_tridiagonal(..., select='v') is not implemented

What it means

jax.scipy.linalg.eigh_tridiagonal does not support select='v' (computing eigenvalues/eigenvectors in a half-open value interval). JAX's implementation relies on static shapes and a bisection counting scheme that needs the count known at trace time, so selecting by value range would require dynamic shape support (per the TODO in the source). Any call with select='v' raises NotImplementedError immediately.

Source

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

  # In the worst case, when the absolute tolerance is eps*lambda_est_max and
  # lambda_est_max = -lambda_est_min, we have to take as many bisection steps
  # as there are bits in the mantissa plus 1.
  # The proof is left as an exercise to the reader.
  max_it = finfo.nmant + 1

  # Determine the indices of the desired eigenvalues, based on select and
  # select_range.
  if select == 'a':
    target_counts = jnp.arange(n, dtype=np.int32)
  elif select == 'i':
    if select_range is None:
      raise ValueError("for select='i', select_range must be specified.")
    if select_range[0] > select_range[1]:
      raise ValueError('Got empty index range in select_range.')
    target_counts = jnp.arange(select_range[0], select_range[1] + 1, dtype=np.int32)
  elif select == 'v':
    # TODO(phawkins): requires dynamic shape support.
    raise NotImplementedError("eigh_tridiagonal(..., select='v') is not "
                              "implemented")
  else:
    raise ValueError("'select must have a value in {'a', 'i', 'v'}.")

  # Run binary search for all desired eigenvalues in parallel, starting from
  # the interval lightly wider than the estimated
  # [lambda_est_min, lambda_est_max].
  fudge = 2.1  # We widen starting interval the Gershgorin interval a bit.
  norm_slack = jnp.array(n, alpha.dtype) * fudge * finfo.eps * t_norm
  lower = lambda_est_min - norm_slack - 2 * fudge * pivmin
  upper = lambda_est_max + norm_slack + fudge * pivmin

  # Pre-broadcast the scalars used in the Sturm sequence for improved
  # performance.
  target_shape = np.shape(target_counts)
  lower = jnp.broadcast_to(lower, shape=target_shape)
  upper = jnp.broadcast_to(upper, shape=target_shape)
  mid = 0.5 * (upper + lower)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use select='i' with an index range instead: compute all eigenvalues once (select='a') and map value thresholds to indices, since JAX returns sorted eigenvalues
2. Use select='a' and post-filter with boolean masks: w, v = eigh_tridiagonal(d, e); mask = (w >= lo) & (w < hi)
3. Fall back to numpy/scipy for this operation (e.g. via host callback or plain scipy outside jit)
4. File/star the upstream issue about dynamic shape support (TODO(phawkins))

Example fix

# before
w, v = eigh_tridiagonal(d, e, select='v', select_range=(lo, hi))
# after
w, v = eigh_tridiagonal(d, e, eigvals_only=False)  # select='a'
mask = (w >= lo) & (w < hi)
w, v = w[mask], v[:, mask]

# No pre-call validation possible (it's an unimplemented feature); check upfront:
def has_value_select(): return False  # JAX eigh_tridiagonal select='v'
if select == 'v':
    w, v = eigh_tridiagonal(d, e)
    mask = (w >= lo) & (w < hi)
    w, v = w[mask], v[:, mask]

try:
    w, v = eigh_tridiagonal(d, e, select='v', select_range=(lo, hi))
except NotImplementedError:
    w, v = eigh_tridiagonal(d, e)
    m = (w >= lo) & (w < hi); w, v = w[m], v[:, m]

When it happens

Trigger: Calling eigh_tridiagonal(d, e, select='v', select_range=(lo, hi)); note select='a' (all) and select='i' (index range) work, only the value-range mode fails.

Common situations: Porting SciPy code that uses scipy.linalg.eigh_tridiagonal(..., select='v') for computing spectrum slices of large tridiagonal Hamiltonians (quantum physics, edge-state computations) to JAX.

Related errors

• Only the type=1 case of eigh is implemented.
• Only the eigvals=None case of eigh is implemented.
• 'select must have a value in {'a', 'i', 'v'}.
• Blocked version is not implemented yet.
• array ref with memory space only works inside of a `jit`.

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