jax-ml/jax · error · ValueError

Unsupported dtype: {dtype}

Error message

Unsupported dtype: {dtype}

What it means

jax/_src/lax/linalg.py:1098 in _eig_gpu_lowering (eig GPU path). Thrown when the dtype of the matrix passed to jax.lax.linalg.eig is not one of the four supported: float32, float64, complex64, complex128. Before lowering to a GPU kernel, JAX classifies the input as real or complex; anything else (e.g. int, float16/bfloat16, or an extended-precision type) reaches the else branch.

Source

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

def _eig_gpu_lowering(ctx, operand, *,
                      compute_left_eigenvectors, compute_right_eigenvectors,
                      enable_eigvec_derivs, implementation, target_name_prefix):
  del enable_eigvec_derivs
  operand_aval, = ctx.avals_in
  batch_dims = operand_aval.shape[:-2]
  n, m = operand_aval.shape[-2:]
  assert n == m

  dtype = operand_aval.dtype
  complex_dtype = np.result_type(dtype, 1j)
  if dtype in (np.float32, np.float64):
    is_real = True
  elif dtype in (np.complex64, np.complex128):
    is_real = False
  else:
    raise ValueError(f"Unsupported dtype: {dtype}")

  have_cusolver_geev = (
      target_name_prefix == "cu"
      and cuda_versions
      and cuda_versions.cusolver_get_version() >= 11701
  )

  if (
      implementation is None and have_cusolver_geev
      and not compute_left_eigenvectors
  ) or implementation == EigImplementation.CUSOLVER:
    if not have_cusolver_geev:
      raise RuntimeError(
          "Nonsymmetric eigendecomposition requires cusolver 11.7.1 or newer"
      )
    if compute_left_eigenvectors:
      raise NotImplementedError(
          "Left eigenvectors are not supported by cusolver")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast the operand to a supported dtype before calling eig: a.astype(jnp.float32) (or complex64/float64/complex128 as appropriate)
2. Check and normalize dtype at model input boundaries, e.g. jax.tree_util.map over params enforcing f32/f64
3. If you truly need float16/bfloat16 results, compute in float32 and cast the outputs back down

Example fix

// before
w, v = jax.lax.linalg.eig(a)  # a is bfloat16
// after
w, v = jax.lax.linalg.eig(a.astype(jnp.float32))

SUPPORTED = {jnp.float32, jnp.float64, np.complex64, np.complex128}
if a.dtype not in SUPPORTED:
    a = a.astype(jnp.float32)

def is_eig_dtype(a: jax.Array) -> bool:
    return a.dtype in (jnp.float32, jnp.float64, jnp.complex64, jnp.complex128)

When it happens

Trigger: Calling jax.lax.linalg.eig (or jax.numpy.linalg.eig) on GPU with an operand dtype outside {float32, float64, complex64, complex128}, e.g. an integer array, float16/bfloat16 weights, or a dtype created via custom casting. Only the GPU lowering raises; CPU may fail elsewhere or differently.

Common situations: Models storing weights in bfloat16 (common in transformers) passed directly to eig; integer matrices from preprocessing; accidental object/odd dtype after mixing NumPy and JAX arrays.

Related errors

• Nonsymmetric eigendecomposition requires cusolver 11.7.1 or
• Left eigenvectors are not supported by cusolver
• cannot cast from `{dtype_name}`
• cannot cast to `{dtype_name}`
• Unsigned integer dtype {aval.dtype} is not supported for dot

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