jax-ml/jax · error · TypeError

Only float32 and float64 inputs are supported as inputs to j

Error message

Only float32 and float64 inputs are supported as inputs to jax.scipy.linalg.eigh_tridiagonal, got {alpha.dtype} and {beta.dtype}

What it means

eigh_tridiagonal only supports float32, float64, complex64, and complex128 (despite the message text mentioning only float32/float64). Inputs with other dtypes — float16, bfloat16, int — raise TypeError.

Source

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

    i, q, count = unrolled_steps((i, q, count))

    # Run the remaining steps of the Sturm sequence using a partially
    # unrolled while loop.
    unroll_cnt = blocksize
    def cond(iqc):
      i, q, count = iqc
      return jnp.less(i, n)
    _, _, count = lax.while_loop(cond, unrolled_steps, (i, q, count))
    return count

  alpha = jnp.asarray(d)
  beta = jnp.asarray(e)
  supported_dtypes = (np.float32, np.float64, np.complex64, np.complex128)
  if alpha.dtype != beta.dtype:
    raise TypeError("diagonal and off-diagonal values must have same dtype, "
                    f"got {alpha.dtype} and {beta.dtype}")
  if alpha.dtype not in supported_dtypes or beta.dtype not in supported_dtypes:
    raise TypeError("Only float32 and float64 inputs are supported as inputs "
                    "to jax.scipy.linalg.eigh_tridiagonal, got "
                    f"{alpha.dtype} and {beta.dtype}")
  n = alpha.shape[0]
  if n <= 1:
    if eigvals_only:
      return jnp.real(alpha)
    else:
      return jnp.real(alpha), jnp.eye(n, dtype=alpha.dtype)

  if dtypes.issubdtype(alpha.dtype, np.complexfloating):
    alpha = jnp.real(alpha)
    beta_sq = jnp.real(beta * jnp.conj(beta))
    beta_abs = jnp.sqrt(beta_sq)
  else:
    beta_abs = jnp.abs(beta)
    beta_sq = jnp.square(beta)

  # Estimate the largest and smallest eigenvalues of T using the Gershgorin

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast inputs: eigh_tridiagonal(d.astype(jnp.float32), e.astype(jnp.float32))
2. For double precision, enable x64 first: jax.config.update('jax_enable_x64', True) and use float64

Example fix

// before
w, v = jax.scipy.linalg.eigh_tridiagonal(d_bf16, e_bf16)
// after
w, v = jax.scipy.linalg.eigh_tridiagonal(d_bf16.astype(jnp.float32), e_bf16.astype(jnp.float32))

_OK = (np.float32, np.float64, np.complex64, np.complex128)
if np.asarray(d).dtype not in _OK: d, e = jnp.asarray(d, np.float32), jnp.asarray(e, np.float32)

_OK = (np.float32, np.float64, np.complex64, np.complex128)
def eigh_tridiag_dtype_ok(d, e): return np.asarray(d).dtype in _OK and np.asarray(e).dtype in _OK

When it happens

Trigger: Calling eigh_tridiagonal with bfloat16/float16 arrays (TPU or mixed-precision defaults), or integer diagonals.

Common situations: Mixed-precision pipelines on TPU; passing Python/NumPy int arrays without conversion to floating point.

Related errors

• diagonal and off-diagonal values must have same dtype, got {
• Unsupported dtype: {dtype}
• A.dtype={A.dtype} is not supported.
• for select='i', select_range must be specified.
• Got empty index range in select_range.

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