jax-ml/jax · error · TypeError

diagonal and off-diagonal values must have same dtype, got {

Error message

diagonal and off-diagonal values must have same dtype, got {alpha.dtype} and {beta.dtype}

What it means

eigh_tridiagonal solves the symmetric tridiagonal eigenproblem from diagonal d and off-diagonal e; the underlying kernels require both arrays to share one dtype. If alpha.dtype != beta.dtype (e.g. float32 d with float64 e), TypeError is raised.

Source

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

        q, count = sturm_step(start + j, q, count)
      return start + unroll_cnt, q, count

    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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast both to the same dtype: eigh_tridiagonal(d.astype(e.dtype), e)
2. Construct both arrays with an explicit common dtype from the start

Example fix

// before
w, v = jax.scipy.linalg.eigh_tridiagonal(d_f32, e_f64)
// after
w, v = jax.scipy.linalg.eigh_tridiagonal(d_f32, e_f64.astype(d_f32.dtype))

if d.dtype != e.dtype: d, e = jnp.asarray(d, e.dtype), jnp.asarray(e, d.dtype)

def same_dtype(d, e): return np.asarray(d).dtype == np.asarray(e).dtype

When it happens

Trigger: Calling jax.scipy.linalg.eigh_tridiagonal(d.astype(jnp.float32), e) where e remained float64, or mixing precisions from different data sources.

Common situations: Loading d and e from datasets stored at different precisions; x64 enabled globally so literals default to float64 while d was constructed float32.

Related errors

• Only float32 and float64 inputs are supported as inputs to j
• 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/c319ef82afa20c6a. Report an issue: GitHub.