jax-ml/jax · error · ValueError

expected A and E to be the same shape, got A.shape={A_arr.sh

Error message

expected A and E to be the same shape, got A.shape={A_arr.shape} E.shape={E_arr.shape}

What it means

expm_frechet differentiates expm via jvp with E as the tangent, which requires A and E to have identical shapes. Mismatched shapes — even both square, e.g. (4,4) vs (8,8), or different batch shapes — raise ValueError.

Source

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

    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:
    *arrs: arrays of at most two dimensions

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make E exactly the same shape as A, e.g. E = jnp.zeros_like(A) for the identity direction
Broadcast explicitly yourself: E = jnp.broadcast_to(E, A.shape)

Example fix

// before
expm_frechet(A, E)  # A: (4,4), E: (8,8)
// after
E = jnp.zeros_like(A); E = E.at[0, 1].set(1.0)
expm_frechet(A, E)

Defensive patterns

Strategy: validation

Validate before calling

if E.shape != A.shape: E = jnp.broadcast_to(E, A.shape)

Type guard

null

Prevention

Always derive E from A via zeros_like/at when possible
Assert A.shape == E.shape in tests

When it happens

Trigger: Calling expm_frechet(A, E) where A.shape != E.shape (different n, different ndim, or different batch dims).

Common situations: Computing a Frechet derivative with respect to a differently-sized perturbation; batched A with unbatched E.

Related errors

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