jax-ml/jax · error · ValueError

Input 'T' must be square.

Error message

Input 'T' must be square.

What it means

jax.scipy.linalg.rsf2csf converts a real Schur form to complex Schur form; the quasi-triangular T matrix must be a (batched) square matrix. The validation rejects arrays with fewer than 2 dimensions or whose last two axes differ in size (shape[-1] != shape[-2]).

Source

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

    [[ 3.76 -2.17  1.38]
     [ 0.   -0.88 -0.35]
     [ 0.    2.37 -0.88]]

    By contrast, the complex form is truly upper-triangular:

    >>> with jnp.printoptions(precision=2, suppress=True):
    ...   print(Tc)
    [[ 3.76+0.j    1.29-0.78j  2.02-0.5j ]
     [ 0.  +0.j   -0.88+0.91j -2.02+0.j  ]
     [ 0.  +0.j    0.  +0.j   -0.88-0.91j]]
  """
  del check_finite  # unused

  T_arr = jnp.asarray(T)
  Z_arr = jnp.asarray(Z)

  if T_arr.ndim < 2 or T_arr.shape[-1] != T_arr.shape[-2]:
    raise ValueError("Input 'T' must be square.")
  if Z_arr.ndim < 2 or Z_arr.shape[-1] != Z_arr.shape[-2]:
    raise ValueError("Input 'Z' must be square.")
  if T_arr.shape[-1] != Z_arr.shape[-1]:
    raise ValueError(f"Input array shapes must match: Z: {Z_arr.shape} vs. T: {T_arr.shape}")

  return jnp_vectorize.vectorize(
      _rsf2csf_2d, signature="(n,n),(n,n)->(n,n),(n,n)")(T_arr, Z_arr)

@overload
def hessenberg(a: ArrayLike, *, calc_q: Literal[False], overwrite_a: bool = False,
               check_finite: bool = True) -> Array: ...

@overload
def hessenberg(a: ArrayLike, *, calc_q: Literal[True], overwrite_a: bool = False,
               check_finite: bool = True) -> tuple[Array, Array]: ...


@jit(static_argnames=('calc_q', 'check_finite', 'overwrite_a'))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Ensure T comes from scipy_schur/schur output or another square matrix source
  2. Check T.shape[-1] == T.shape[-2] and T.ndim >= 2 before calling
  3. Fix upstream slicing that dropped an axis (e.g. T[0] vs T[:, 0])

Example fix

# before
Tc, Zc = rsf2csf(t_vec, z)  # t_vec shape (n,)
# after
Tc, Zc = rsf2csf(t_vec.reshape(-1, 1), z)
Defensive patterns

Strategy: validation

Validate before calling

T = jnp.asarray(T)
if T.ndim < 2 or T.shape[-1] != T.shape[-2]:
    raise ValueError(f"T must be square, got {T.shape}")
Tc, Zc = rsf2csf(T, Z)

Type guard

def is_square_matrix(x) -> bool:
    x = jnp.asarray(x)
    return x.ndim >= 2 and x.shape[-1] == x.shape[-2]

Try / catch

try:
    rsf2csf(T, Z)
except ValueError as e:
    if "must be square" in str(e):
        raise ValueError(f'bad Schur factors: {e}') from e
    raise

Prevention

When it happens

Trigger: Passing a 1-D array or a ragged/non-square last-two-dims array as T, e.g. shape (3, 4) or (2, 2, 3, 4).

Common situations: Feeding the output of schur() that was reshaped/sliced incorrectly; passing a vector of eigenvalues instead of the Schur factor; batched pipelines with a malformed batch member.

Related errors


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