jax-ml/jax · error · ValueError
Input array shapes must match: Z: {Z_arr.shape} vs. T: {T_ar
Error message
Input array shapes must match: Z: {Z_arr.shape} vs. T: {T_arr.shape} What it means
rsf2csf requires T and Z to describe the same problem, so their trailing dimensions must match exactly (T.shape[-1] == Z.shape[-1]). Mismatched sizes (e.g. T is 4x4 and Z is 3x3) raise ValueError with both shapes in the message.
Source
Thrown at jax/_src/scipy/linalg.py:2333
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]:
"""Compute the Hessenberg form of the matrix
View on GitHub (pinned to 1e1c6a8fc0)
Solutions
- Recompute both T and Z from the same schur(a, calc_q=True) call
- Assert T.shape[-1] == Z.shape[-1] before calling
- If batching, ensure both arrays have matching batch dims too
Example fix
# before T, _ = schur(a, calc_q=False) _, Z = schur(b, calc_q=True) # different size than a Tc, Zc = rsf2csf(T, Z) # after T, Z = schur(a, calc_q=True) Tc, Zc = rsf2csf(T, Z)
Defensive patterns
Strategy: validation
Validate before calling
T, Z = jnp.asarray(T), jnp.asarray(Z)
if T.shape[-1] != Z.shape[-1]:
raise ValueError(f'T/Z size mismatch: {T.shape} vs {Z.shape}')
Tc, Zc = rsf2csf(T, Z) Try / catch
try:
rsf2csf(T, Z)
except ValueError as e:
if 'shapes must match' in str(e):
T, Z = schur(A, calc_q=True); rsf2csf(T, Z)
else: raise Prevention
- Never cache one factor across problems of changing size
- Pair T and Z atomically (return both from one function)
- Assert matching trailing dims in batched pipelines
When it happens
Trigger: Calling rsf2csf(T, Z) where T and Z come from Schur decompositions of different-sized matrices, or one was sliced/reshaped independently of the other.
Common situations: Caching or reusing factors across iterations where matrix size changed; mixing batch entries from different runs; off-by-one slicing of batched factors.
Related errors
- unexpected JAX type (e.g. shape/dtype) for argument to VJP f
- cotangent type does not match function output, expected {out
- Mismatched number of outputs from callback. Expected: {}, Ac
- Incorrect output shape for return value #{i}: Expected: {out
- Length of sharding.spec ({len(out_s.spec)}) must be equal to
AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27).
Data as JSON: /api/errors/f165dffe9edd28f7.
Report an issue: GitHub.