{"record":{"id":"0278add0ea2a153a","repo":"jax-ml/jax","slug":"expected-a-to-be-a-batched-square-matrix-got-a","errorCode":null,"errorMessage":"Expected A to be a (batched) square matrix, got {A.shape=}.","messagePattern":"Expected A to be a \\(batched\\) square matrix, got (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/linalg.py","lineNumber":1424,"sourceCode":"    >>> jnp.allclose(jax.scipy.linalg.expm(A+B),\n    ...              jax.scipy.linalg.expm(A) @ jax.scipy.linalg.expm(B),\n    ...              rtol=0.0001)\n    Array(True, dtype=bool)\n\n    If a matrix ``X`` is invertible, then\n    ``expm(X @ A @ inv(X)) = X @ expm(A) @ inv(X)``\n\n    >>> X = jnp.array([[3, 1],\n    ...                [2, 5]])\n    >>> X_inv = jax.scipy.linalg.inv(X)\n    >>> jnp.allclose(jax.scipy.linalg.expm(X @ A @ X_inv),\n    ...              X @ jax.scipy.linalg.expm(A) @ X_inv)\n    Array(True, dtype=bool)\n  \"\"\"\n  A, = promote_dtypes_inexact(A)\n\n  if A.ndim < 2 or A.shape[-1] != A.shape[-2]:\n    raise ValueError(f\"Expected A to be a (batched) square matrix, got {A.shape=}.\")\n\n  if A.ndim > 2:\n    return jnp_vectorize.vectorize(\n      partial(expm, upper_triangular=upper_triangular, max_squarings=max_squarings),\n      signature=\"(n,n)->(n,n)\")(A)\n\n  P, Q, n_squarings = _calc_P_Q(jnp.asarray(A))\n\n  def _nan(args):\n    A, *_ = args\n    return jnp.full_like(A, np.nan)\n\n  def _compute(args):\n    A, P, Q = args\n    R = _solve_P_Q(P, Q, upper_triangular)\n    R = _squaring(R, n_squarings, max_squarings)\n    return R\n","sourceCodeStart":1406,"sourceCodeEnd":1442,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/linalg.py#L1406-L1442","documentation":"jax.scipy.linalg.expm computes the matrix exponential via scaling-and-squaring Padé approximation, which is only defined for square matrices. It requires A.ndim >= 2 and A.shape[-1] == A.shape[-2] (batched square); otherwise ValueError.","triggerScenarios":"Calling jax.scipy.linalg.expm on a 1-D array (treated as a vector), a non-square (m, n) matrix with m != n, or a batch whose leaf matrices are non-square.","commonSituations":"Passing a vector of rates expecting elementwise expm (should use jnp.exp); a data-shape bug (transposed batch or off-by-one reshape) silently producing rectangular matrices.","solutions":["Use jnp.exp for elementwise exponentials of vectors","Reshape/fix the input to square matrices, e.g. (n, n) or batch (..., n, n)","Assert squareness before calling: assert A.ndim >= 2 and A.shape[-1] == A.shape[-2]"],"exampleFix":"// before\nout = jax.scipy.linalg.expm(jnp.array([1.0, 2.0]))\n// after\nout = jnp.exp(jnp.array([1.0, 2.0]))","handlingStrategy":"validation","validationCode":"assert A.ndim >= 2 and A.shape[-1] == A.shape[-2], 'expm needs square matrices'","typeGuard":"def is_square_batched(A): return A.ndim >= 2 and A.shape[-1] == A.shape[-2]","tryCatchPattern":null,"preventionTips":["Use jnp.exp for elementwise exponentials","Add squareness asserts in matrix-heavy pipelines"],"tags":["jax","linalg","matrix-exponential","shape-mismatch"],"backgroundTag":"matrix-dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}