jax-ml/jax · error · ValueError

Input 'Z' must be square.

Error message

Input 'Z' must be square.

What it means

In rsf2csf, the unitary transformation matrix Z must also be square (ndim >= 2 and shape[-1] == shape[-2]), matching the Schur factor it accompanies. Non-square or lower-dimensional Z fails this check right after the T check.

Source

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

     [ 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'))
def hessenberg(a: ArrayLike, *, calc_q: bool = False, overwrite_a: bool = False,
               check_finite: bool = True) -> Array | tuple[Array, Array]:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass the full square Z from schur(a)
2. Verify Z.ndim >= 2 and Z.shape[-1] == Z.shape[-2] before calling
3. Regenerate Z with calc_q=True in schur rather than reconstructing a reduced basis

Example fix

# before
T, Z = schur(a, calc_q=False)
Tc, Zc = rsf2csf(T, Z)  # Z is None -> fails earlier; or wrong Z shape
# after
T, Z = schur(a, calc_q=True)
Tc, Zc = rsf2csf(T, Z)

Z = jnp.asarray(Z)
assert Z.ndim >= 2 and Z.shape[-1] == Z.shape[-2], f'Z must be square, got {Z.shape}'
Tc, Zc = rsf2csf(T, Z)

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

try:
    rsf2csf(T, Z)
except ValueError as e:
    if "'Z' must be square" in str(e):
        T, Z = schur(A, calc_q=True); rsf2csf(T, Z)
    else: raise

When it happens

Trigger: Passing a 1-D or rectangular Z (e.g. the Q from a reduced QR, or a flattened matrix) to rsf2csf.

Common situations: Using a compacted/thin factor from another decomposition instead of the full Schur vectors; accidental reshape or axis drop in batched code.

Related errors

• Input 'T' must be square.
• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch
• conv_general_dilated rhs output feature dimension size must
• conv_general_dilated window and window_strides must have the

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