{"record":{"id":"8bfe7fb3e91de5c2","repo":"jax-ml/jax","slug":"expected-a-to-be-a-square-matrix","errorCode":null,"errorMessage":"expected A to be a square matrix","messagePattern":"expected A to be a square matrix","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/scipy/linalg.py","lineNumber":1449,"sourceCode":"  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\n  R = lax.cond(n_squarings > max_squarings, _nan, _compute, (A, P, Q))\n  return R\n\n@jit\ndef _calc_P_Q(A: Array) -> tuple[Array, Array, Array]:\n  if A.ndim != 2 or A.shape[0] != A.shape[1]:\n    raise ValueError('expected A to be a square matrix')\n  A_L1 = jnp_linalg.norm(A,1)\n  n_squarings: Array\n  U: Array\n  V: Array\n  if A.dtype == 'float64' or A.dtype == 'complex128':\n   maxnorm = 5.371920351148152\n   n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))\n   A = A / 2 ** n_squarings.astype(A.dtype)\n   conds = jnp.array([1.495585217958292e-002, 2.539398330063230e-001,\n                      9.504178996162932e-001, 2.097847961257068e+000],\n                      dtype=A_L1.dtype)\n   idx = jnp.digitize(A_L1, conds)\n   U, V = lax.switch(idx, [_pade3, _pade5, _pade7, _pade9, _pade13], A)\n  elif A.dtype == 'float32' or A.dtype == 'complex64':\n    maxnorm = 3.925724783138660\n    n_squarings = jnp.maximum(0, jnp.floor(jnp.log2(A_L1 / maxnorm)))\n    A = A / 2 ** n_squarings.astype(A.dtype)\n    conds = jnp.array([4.258730016922831e-001, 1.880152677804762e+000],","sourceCodeStart":1431,"sourceCodeEnd":1467,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/scipy/linalg.py#L1431-L1467","documentation":"The internal helper _calc_P_Q (used by expm for the Padé approximation) hard-requires a single 2-D square matrix. Because it runs under jit as a traced function, this check only fires in practice when the function is called eagerly (e.g. from a test or direct call) with a non-square matrix.","triggerScenarios":"Calling jax._src.scipy.linalg._calc_P_Q directly with a non-square or batched array; or expm hitting this path without its public wrapper's shape guard (custom/vmap code paths).","commonSituations":"Vendoring or reusing JAX internals; calling expm on data that bypassed the public validation, e.g. via custom wrappers stripping the check.","solutions":["Don't call _calc_P_Q directly; use jax.scipy.linalg.expm which handles batching via vectorize","Ensure the input is a single square (n, n) matrix","For batches, rely on expm's built-in vmap over leading dims"],"exampleFix":"// before\nP, Q, ns = jax._src.scipy.linalg._calc_P_Q(A_rect)  # (3, 5)\n// after\nP, Q, ns = jax._src.scipy.linalg._calc_P_Q(A_square)  # (n, n); or use expm","handlingStrategy":"validation","validationCode":"assert A.ndim == 2 and A.shape[0] == A.shape[1]","typeGuard":"null","tryCatchPattern":null,"preventionTips":["Avoid private jax._src APIs; they lack the public wrappers' guards","Let expm handle batching via its internal vectorize"],"tags":["jax","linalg","internal-api","shape-mismatch"],"backgroundTag":"matrix-dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}