{"record":{"id":"5c70176f77fd052a","repo":"jax-ml/jax","slug":"expected-a-and-e-to-be-the-same-shape-got-a-shape","errorCode":null,"errorMessage":"expected A and E to be the same shape, got A.shape={A_arr.shape} E.shape={E_arr.shape}","messagePattern":"expected A and E to be the same shape, got A\\.shape=(.+?) E\\.shape=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/linalg.py","lineNumber":1618,"sourceCode":"    This can be equivalently computed using JAX's automatic differentiation methods;\n    here we'll compute the derivative of :func:`~jax.scipy.linalg.expm` in the\n    direction of ``E`` using :func:`jax.jvp`, and find the same results:\n\n    >>> expmA2, expm_frechet_AE2 = jax.jvp(jax.scipy.linalg.expm, (A,), (E,))\n    >>> jnp.allclose(expmA, expmA2)\n    Array(True, dtype=bool)\n    >>> jnp.allclose(expm_frechet_AE, expm_frechet_AE2)\n    Array(True, dtype=bool)\n  \"\"\"\n  del method  # unused\n  A_arr = jnp.asarray(A)\n  E_arr = jnp.asarray(E)\n  if A_arr.ndim < 2 or A_arr.shape[-2] != A_arr.shape[1]:\n    raise ValueError(f'expected A to be a (batched) square matrix, got A.shape={A_arr.shape}')\n  if E_arr.ndim < 2 or E_arr.shape[-2] != E_arr.shape[-1]:\n    raise ValueError(f'expected E to be a (batched) square matrix, got E.shape={E_arr.shape}')\n  if A_arr.shape != E_arr.shape:\n    raise ValueError('expected A and E to be the same shape, got '\n                     f'A.shape={A_arr.shape} E.shape={E_arr.shape}')\n  bound_fun = partial(expm, upper_triangular=False, max_squarings=16)\n  expm_A, expm_frechet_AE = jvp(bound_fun, (A_arr,), (E_arr,))\n  if compute_expm:\n    return expm_A, expm_frechet_AE\n  else:\n    return expm_frechet_AE\n\n\n@jit\ndef block_diag(*arrs: ArrayLike) -> Array:\n  \"\"\"Create a block diagonal matrix from input arrays.\n\n  JAX implementation of :func:`scipy.linalg.block_diag`.\n\n  Args:\n    *arrs: arrays of at most two dimensions\n","sourceCodeStart":1600,"sourceCodeEnd":1636,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/linalg.py#L1600-L1636","documentation":"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.","triggerScenarios":"Calling expm_frechet(A, E) where A.shape != E.shape (different n, different ndim, or different batch dims).","commonSituations":"Computing a Frechet derivative with respect to a differently-sized perturbation; batched A with unbatched E.","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)"],"exampleFix":"// before\nexpm_frechet(A, E)  # A: (4,4), E: (8,8)\n// after\nE = jnp.zeros_like(A); E = E.at[0, 1].set(1.0)\nexpm_frechet(A, E)","handlingStrategy":"validation","validationCode":"if E.shape != A.shape: E = jnp.broadcast_to(E, A.shape)","typeGuard":"null","tryCatchPattern":null,"preventionTips":["Always derive E from A via zeros_like/at when possible","Assert A.shape == E.shape in tests"],"tags":["jax","linalg","frechet-derivative","shape-mismatch"],"backgroundTag":"matrix-dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}