jax-ml/jax · error · ValueError

expected E to be a (batched) square matrix, got E.shape={E_a

Error message

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

What it means

expm_frechet's direction matrix E must itself be a (batched) square matrix (ndim >= 2, last two dims equal), matching the structure of A. A non-square or vector E raises ValueError before the jvp is computed.

Source

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

    >>> 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`.

  Args:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Reshape E to (n, n): E = E.reshape(n, n)
Ensure E has the same batch shape structure as A

Example fix

// before
expm_frechet(A, e_vec)  # e_vec.shape == (n,)
// after
expm_frechet(A, e_vec.reshape(n, n))

Defensive patterns

Strategy: validation

Validate before calling

assert E.ndim >= 2 and E.shape[-2] == E.shape[-1], 'E must be square'

Type guard

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

Prevention

Reshape flattened direction vectors to (n, n) before calling
Keep A and E construction in one helper so shapes stay consistent

When it happens

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

Common situations: Passing a perturbation direction as a flattened vector instead of a matrix; broadcasting bugs making E rectangular.

Related errors

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