jax-ml/jax · error · NotImplementedError

Jacobi implementation is not supported on CPU

Error message

Jacobi implementation is not supported on CPU

What it means

jax/_src/lax/linalg.py:1295 in _eigh_cpu_gpu_lowering. EighImplementation.JACOBI is implemented only for GPU; selecting it while the target is CPU raises NotImplementedError. On CPU the LAPACK syevd/heevd path is used instead.

Source

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

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:
      algo_int = 2 if algorithm == EighImplementation.JACOBI else 1
    kwargs = {"lower": lower, "algorithm": np.uint8(algo_int)}

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Remove the explicit algorithm argument to use the LAPACK default on CPU
  2. Conditionally select algorithm based on jax.default_backend()
  3. Ensure the GPU is actually visible (nvidia-smi, jax.devices()) if Jacobi was intended

Example fix

// before
evals, evecs = jax.lax.linalg.eigh(a, algorithm=lax.linalg.EighImplementation.JACOBI)  # on CPU
// after
algo = lax.linalg.EighImplementation.JACOBI if jax.default_backend() == 'gpu' else None
evals, evecs = jax.lax.linalg.eigh(a, algorithm=algo)
Defensive patterns

Strategy: validation

Validate before calling

algo = (lax.linalg.EighImplementation.JACOBI
         if jax.default_backend() == 'gpu' else None)

Prevention

When it happens

Trigger: Calling jax.lax.linalg.eigh(a, algorithm=EighImplementation.JACOBI) with jax backend cpu (e.g. JAX_PLATFORMS=cpu or no GPU present).

Common situations: GPU-authored config reused in CPU-only CI or a laptop; jax.config defaults changed; debugging numerics by forcing Jacobi for its better accuracy on GPU then running tests on CPU.

Related errors

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