{"record":{"id":"bc3d53a373f3f4e8","repo":"jax-ml/jax","slug":"expected-a-to-be-a-batched-square-matrix-got-a-bc3d53","errorCode":null,"errorMessage":"expected A to be a (batched) square matrix, got A.shape={A_arr.shape}","messagePattern":"expected A to be a \\(batched\\) square matrix, got A\\.shape=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/linalg.py","lineNumber":1614,"sourceCode":"    >>> A = jax.random.normal(key1, (3, 3))\n    >>> E = jax.random.normal(key2, (3, 3))\n    >>> expmA, expm_frechet_AE = jax.scipy.linalg.expm_frechet(A, E)\n\n    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`.","sourceCodeStart":1596,"sourceCodeEnd":1632,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/linalg.py#L1596-L1632","documentation":"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.","triggerScenarios":"Calling expm_frechet(A, E) with A of shape (m, n), m != n, or with a vector A.","commonSituations":"Condition-number estimation pipelines (expm_cond) feeding misshapen matrices; upstream reshape bugs producing rectangular A.","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]"],"exampleFix":"// before\nexpm_frechet(jnp.ones((3, 4)), E)\n// after\nexpm_frechet(jnp.ones((4, 4)), E)","handlingStrategy":"validation","validationCode":"assert A.ndim >= 2 and A.shape[-2] == A.shape[-1]","typeGuard":"def is_square_batched(A): return A.ndim >= 2 and A.shape[-2] == A.shape[-1]","tryCatchPattern":null,"preventionTips":["Validate both A and E shapes before expm_frechet","Use zeros_like(A) to build direction matrices with guaranteed shape"],"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"}