jax-ml/jax · error · ValueError

Expected A to be a (batched) square matrix, got {A.shape=}.

Error message

Expected A to be a (batched) square matrix, got {A.shape=}.

What it means

jax.scipy.linalg.expm computes the matrix exponential via scaling-and-squaring Padé approximation, which is only defined for square matrices. It requires A.ndim >= 2 and A.shape[-1] == A.shape[-2] (batched square); otherwise ValueError.

Source

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

    >>> jnp.allclose(jax.scipy.linalg.expm(A+B),
    ...              jax.scipy.linalg.expm(A) @ jax.scipy.linalg.expm(B),
    ...              rtol=0.0001)
    Array(True, dtype=bool)

    If a matrix ``X`` is invertible, then
    ``expm(X @ A @ inv(X)) = X @ expm(A) @ inv(X)``

    >>> X = jnp.array([[3, 1],
    ...                [2, 5]])
    >>> X_inv = jax.scipy.linalg.inv(X)
    >>> jnp.allclose(jax.scipy.linalg.expm(X @ A @ X_inv),
    ...              X @ jax.scipy.linalg.expm(A) @ X_inv)
    Array(True, dtype=bool)
  """
  A, = promote_dtypes_inexact(A)

  if A.ndim < 2 or A.shape[-1] != A.shape[-2]:
    raise ValueError(f"Expected A to be a (batched) square matrix, got {A.shape=}.")

  if A.ndim > 2:
    return jnp_vectorize.vectorize(
      partial(expm, upper_triangular=upper_triangular, max_squarings=max_squarings),
      signature="(n,n)->(n,n)")(A)

  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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Use jnp.exp for elementwise exponentials of vectors
Reshape/fix the input to square matrices, e.g. (n, n) or batch (..., n, n)
Assert squareness before calling: assert A.ndim >= 2 and A.shape[-1] == A.shape[-2]

Example fix

// before
out = jax.scipy.linalg.expm(jnp.array([1.0, 2.0]))
// after
out = jnp.exp(jnp.array([1.0, 2.0]))

Defensive patterns

Strategy: validation

Validate before calling

assert A.ndim >= 2 and A.shape[-1] == A.shape[-2], 'expm needs square matrices'

Type guard

def is_square_batched(A): return A.ndim >= 2 and A.shape[-1] == A.shape[-2]

Prevention

Use jnp.exp for elementwise exponentials
Add squareness asserts in matrix-heavy pipelines

When it happens

Trigger: Calling jax.scipy.linalg.expm on a 1-D array (treated as a vector), a non-square (m, n) matrix with m != n, or a batch whose leaf matrices are non-square.

Common situations: Passing a vector of rates expecting elementwise expm (should use jnp.exp); a data-shape bug (transposed batch or off-by-one reshape) silently producing rectangular matrices.

Related errors

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