jax-ml/jax · error · ValueError

Argument to Hessenberg reduction must have shape [..., n, n]

Error message

Argument to Hessenberg reduction must have shape [..., n, n], got shape {shape}

What it means

jax/_src/lax/linalg.py:1381 in _hessenberg_shape_rule. jax.lax.linalg.hessenberg reduces a square matrix to upper Hessenberg form; the shape rule requires the last two dims equal ([..., n, n]). A rectangular operand raises ValueError with the observed shape.

Source

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

                precision=lax.Precision.HIGHEST)
  vdag_adot_v = dot(dot(_H(v), a_dot), v)
  dv = dot(v, Fmat * vdag_adot_v)
  dw = _extract_diagonal(vdag_adot_v.real)
  return (v, w_real), (dv, dw)


eigh_p = linalg_primitive(
    _eigh_dtype_rule, (_float | _complex,), (2,), _eigh_shape_rule, "eigh",
    multiple_results=True)
ad.primitive_jvps[eigh_p] = _eigh_jvp_rule
register_cpu_gpu_lowering(eigh_p, _eigh_cpu_gpu_lowering)


# Hessenberg reduction

def _hessenberg_shape_rule(shape, **_):
  if shape[0] != shape[-1]:
    raise ValueError(
        "Argument to Hessenberg reduction must have shape [..., n, n], "
        f"got shape {shape}"
    )
  return shape, shape[:-2] + (shape[-1] - 1,)


def _hessenberg_dtype_rule(dtype, **_):
  return dtype, dtype


def _hessenberg_cpu_lowering(ctx, a):
  a_aval, = ctx.avals_in
  batch_dims = a_aval.shape[:-2]
  n = a_aval.shape[-1]
  if not core.is_constant_dim(n):
    raise ValueError("hessenberg requires the last dimension of a to be "
                     f"constant, got a.shape of {a.shape}.")
  target_name = lapack.prepare_lapack_call("gehrd_ffi", a_aval.dtype)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the operand is square: assert a.shape[-2] == a.shape[-1]
2. If you meant to tridiagonalize/bidiagonalize a rectangular matrix, use qr or svd-based preprocessing instead
3. Fix upstream shape construction (e.g. use A @ A.T or correct slicing)

Example fix

// before
h, q = jax.lax.linalg.hessenberg(A)  # A: (m, k), m != k
// after
assert A.shape[-2] == A.shape[-1]
h, q = jax.lax.linalg.hessenberg(A)

assert a.ndim >= 2 and a.shape[-2] == a.shape[-1]

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.hessenberg on an array whose last two dimensions differ, e.g. (m, k) with m != k, or a batch of non-square matrices.

Common situations: Applying Hessenberg reduction as a preprocessing step for eigensolvers on data matrices instead of square operators; transposition/slicing bugs producing off-by-one shapes.

Related errors

• Argument to symmetric eigendecomposition must have shape [..
• 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/b06cefc435f5ba43. Report an issue: GitHub.