jax-ml/jax · error · ValueError

Argument to symmetric eigendecomposition must have shape [..

Error message

Argument to symmetric eigendecomposition must have shape [..., n, n], got shape {shape}

What it means

jax/_src/lax/linalg.py:1268 in _eigh_shape_rule. eigh (symmetric/Hermitian eigendecomposition) only accepts square matrices: the last two dimensions of the operand must be equal (shape [..., n, n]). The rule compares shape[-2] != shape[-1] and reports the (unbatched) shape it sees.

Source

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

    tangent_out.append(dvl)
  if compute_right_eigenvectors:
    primal_out.append(vr)
    tangent_out.append(dvr)
  return primal_out, tangent_out

eig_p = linalg_primitive(
    _eig_dtype_rule, (_float | _complex,), (2,), _eig_shape_rule, "eig",
    multiple_results=True)
ad.primitive_jvps[eig_p] = eig_jvp_rule
mlir.register_lowering(eig_p, _eig_cpu_lowering, platform="cpu")
register_cpu_gpu_lowering(eig_p, _eig_gpu_lowering, ("cuda", "rocm", "oneapi"))


# Symmetric/Hermitian eigendecomposition

def _eigh_shape_rule(shape, *, subset_by_index, **_):
  if shape[0] != shape[-1]:
    raise ValueError(
        "Argument to symmetric eigendecomposition must have shape [..., n, n], "
        f"got shape {shape}"
    )
  n = shape[0]
  d = (n if subset_by_index is None else
       subset_by_index[1] - subset_by_index[0])
  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Verify a.shape[-2] == a.shape[-1] and construct the symmetric matrix properly (e.g. cov = X.T @ X or jnp.cov)
2. If you wanted singular values of a rectangular matrix, use jnp.linalg.svd instead
3. Add an assert before the call during development to catch shape mistakes early

Example fix

// before
evals, evecs = jnp.linalg.eigh(X)  # X: (N, d) raw data
// after
cov = X.T @ X / (X.shape[0] - 1)  # (d, d)
evals, evecs = jnp.linalg.eigh(cov)

assert a.shape[-2] == a.shape[-1], f'eigh needs square, got {a.shape}'

def is_square(a: jax.Array) -> bool:
    return a.ndim >= 2 and a.shape[-2] == a.shape[-1]

When it happens

Trigger: Calling jax.lax.linalg.eigh / jnp.linalg.eigh on a non-square array, e.g. shape (m, k) with m != k, or a batch of rectangular matrices; often from passing a covariance/correlation matrix that was miscomputed (e.g. X @ X.T with wrong axis) or flattening batch dims incorrectly.

Common situations: Feeding raw data (N x d) instead of the Gram/covariance (d x d); transposition bugs; ragged batches padded inconsistently; applying eigh where svd was intended.

Related errors

• Argument to Hessenberg reduction must have shape [..., n, n]
• The first argument to householder_product must have at least
• The second argument to householder_product must not have mor
• ormqr with left=True expects c to have the same number of ro
• ormqr with left=False expects c to have the same number of c

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