jax-ml/jax · error · NotImplementedError

QDWH implementation is only supported on TPU

Error message

QDWH implementation is only supported on TPU

What it means

jax/_src/lax/linalg.py:1293 in _eigh_cpu_gpu_lowering. EighImplementation.QDWH (a QR-based Dynamically Weighted Halley iteration used on TPU) has no CPU/GPU implementation. Explicitly selecting algorithm=EighImplementation.QDWH while executing on CPU or GPU raises NotImplementedError.

Source

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

  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"
    # Use Jacobi (algorithm=2) if requested, otherwise use QR (algorithm=1)
    if algorithm is None:
      algo_int = 0
    else:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Drop the explicit algorithm and use the default per-backend implementation
  2. Select a supported algorithm for CPU/GPU (e.g. EighImplementation.LAPACK-ish default, or JACOBI on GPU)
  3. Run that section on TPU via jax.device_put with tpu backend if QDWH semantics are required

Example fix

// before
evals, evecs = jax.lax.linalg.eigh(a, algorithm=lax.linalg.EighImplementation.QDWH)
// after (on CPU/GPU)
evals, evecs = jax.lax.linalg.eigh(a)
Defensive patterns

Strategy: validation

Validate before calling

if jax.default_backend() != 'tpu':
    algorithm = None  # never force QDWH off-TPU

Prevention

When it happens

Trigger: Calling jax.lax.linalg.eigh(a, algorithm=EighImplementation.QDWH) with backend cpu or gpu. Auto selection (algorithm=None) never picks QDWH off-TPU, so this only occurs with explicit configuration.

Common situations: Config copied from a TPU training pipeline to a local GPU/CPU dev machine; experimenting with algorithms for numerical accuracy; library code that hardcodes QDWH for determinism across TPU replicas.

Related errors

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