jax-ml/jax · error · ValueError

expected A to be a (batched) square matrix, got A.shape={A_a

Error message

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

What it means

jax.scipy.linalg.expm_frechet computes the Frechet derivative of expm via jvp; it requires A to be a (batched) square matrix (ndim >= 2 and last two dims equal). Non-square or 1-D A raises ValueError.

Source

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

    >>> A = jax.random.normal(key1, (3, 3))
    >>> E = jax.random.normal(key2, (3, 3))
    >>> expmA, expm_frechet_AE = jax.scipy.linalg.expm_frechet(A, E)

    This can be equivalently computed using JAX's automatic differentiation methods;
    here we'll compute the derivative of :func:`~jax.scipy.linalg.expm` in the
    direction of ``E`` using :func:`jax.jvp`, and find the same results:

    >>> expmA2, expm_frechet_AE2 = jax.jvp(jax.scipy.linalg.expm, (A,), (E,))
    >>> jnp.allclose(expmA, expmA2)
    Array(True, dtype=bool)
    >>> jnp.allclose(expm_frechet_AE, expm_frechet_AE2)
    Array(True, dtype=bool)
  """
  del method  # unused
  A_arr = jnp.asarray(A)
  E_arr = jnp.asarray(E)
  if A_arr.ndim < 2 or A_arr.shape[-2] != A_arr.shape[1]:
    raise ValueError(f'expected A to be a (batched) square matrix, got A.shape={A_arr.shape}')
  if E_arr.ndim < 2 or E_arr.shape[-2] != E_arr.shape[-1]:
    raise ValueError(f'expected E to be a (batched) square matrix, got E.shape={E_arr.shape}')
  if A_arr.shape != E_arr.shape:
    raise ValueError('expected A and E to be the same shape, got '
                     f'A.shape={A_arr.shape} E.shape={E_arr.shape}')
  bound_fun = partial(expm, upper_triangular=False, max_squarings=16)
  expm_A, expm_frechet_AE = jvp(bound_fun, (A_arr,), (E_arr,))
  if compute_expm:
    return expm_A, expm_frechet_AE
  else:
    return expm_frechet_AE


@jit
def block_diag(*arrs: ArrayLike) -> Array:
  """Create a block diagonal matrix from input arrays.

  JAX implementation of :func:`scipy.linalg.block_diag`.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Fix A to be square: (n, n) or batch (..., n, n)
Add an assert before the call: assert A.ndim >= 2 and A.shape[-2] == A.shape[-1]

Example fix

// before
expm_frechet(jnp.ones((3, 4)), E)
// after
expm_frechet(jnp.ones((4, 4)), E)

Defensive patterns

Strategy: validation

Validate before calling

assert A.ndim >= 2 and A.shape[-2] == A.shape[-1]

Type guard

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

Prevention

Validate both A and E shapes before expm_frechet
Use zeros_like(A) to build direction matrices with guaranteed shape

When it happens

Trigger: Calling expm_frechet(A, E) with A of shape (m, n), m != n, or with a vector A.

Common situations: Condition-number estimation pipelines (expm_cond) feeding misshapen matrices; upstream reshape bugs producing rectangular A.

Related errors

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