jax-ml/jax · error · ValueError

expected A to be a square matrix

Error message

expected A to be a square matrix

What it means

The internal helper _calc_P_Q (used by expm for the Padé approximation) hard-requires a single 2-D square matrix. Because it runs under jit as a traced function, this check only fires in practice when the function is called eagerly (e.g. from a test or direct call) with a non-square matrix.

Source

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

  P, Q, n_squarings = _calc_P_Q(jnp.asarray(A))

  def _nan(args):
    A, *_ = args
    return jnp.full_like(A, np.nan)

  def _compute(args):
    A, P, Q = args
    R = _solve_P_Q(P, Q, upper_triangular)
    R = _squaring(R, n_squarings, max_squarings)
    return R

  R = lax.cond(n_squarings > max_squarings, _nan, _compute, (A, P, Q))
  return R

@jit
def _calc_P_Q(A: Array) -> tuple[Array, Array, Array]:
  if A.ndim != 2 or A.shape[0] != A.shape[1]:
    raise ValueError('expected A to be a square matrix')
  A_L1 = jnp_linalg.norm(A,1)
  n_squarings: Array
  U: Array
  V: Array
  if A.dtype == 'float64' or A.dtype == 'complex128':
   maxnorm = 5.371920351148152
   n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))
   A = A / 2 ** n_squarings.astype(A.dtype)
   conds = jnp.array([1.495585217958292e-002, 2.539398330063230e-001,
                      9.504178996162932e-001, 2.097847961257068e+000],
                      dtype=A_L1.dtype)
   idx = jnp.digitize(A_L1, conds)
   U, V = lax.switch(idx, [_pade3, _pade5, _pade7, _pade9, _pade13], A)
  elif A.dtype == 'float32' or A.dtype == 'complex64':
    maxnorm = 3.925724783138660
    n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))
    A = A / 2 ** n_squarings.astype(A.dtype)
    conds = jnp.array([4.258730016922831e-001, 1.880152677804762e+000],

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Don't call _calc_P_Q directly; use jax.scipy.linalg.expm which handles batching via vectorize
2. Ensure the input is a single square (n, n) matrix
3. For batches, rely on expm's built-in vmap over leading dims

Example fix

// before
P, Q, ns = jax._src.scipy.linalg._calc_P_Q(A_rect)  # (3, 5)
// after
P, Q, ns = jax._src.scipy.linalg._calc_P_Q(A_square)  # (n, n); or use expm

assert A.ndim == 2 and A.shape[0] == A.shape[1]

null

When it happens

Trigger: Calling jax._src.scipy.linalg._calc_P_Q directly with a non-square or batched array; or expm hitting this path without its public wrapper's shape guard (custom/vmap code paths).

Common situations: Vendoring or reusing JAX internals; calling expm on data that bypassed the public validation, e.g. via custom wrappers stripping the check.

Related errors

• multi_dot: last dimension of each array must match first dim
• Array shapes are not compatible for Q @ c operation: a has s
• Array shapes are not compatible for c @ Q operation: a has s
• Expected A to be a (batched) square matrix, got {A.shape=}.
• expected A to be a (batched) square matrix, got A.shape={A_a

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