{"record":{"id":"a66bbf3c8eddc239","repo":"jax-ml/jax","slug":"last-two-dimensions-of-lu-decomposition-must-be-eq","errorCode":null,"errorMessage":"last two dimensions of LU decomposition must be equal, got shape {lu.shape}","messagePattern":"last two dimensions of LU decomposition must be equal, got shape (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/linalg.py","lineNumber":1830,"sourceCode":"    x = triangular_solve(lu, x, left_side=True, lower=True, unit_diagonal=True)\n    x = triangular_solve(lu, x, left_side=True, lower=False)\n  elif trans == 1 or trans == 2:\n    conj = trans == 2\n    x = triangular_solve(lu, x, left_side=True, lower=False, transpose_a=True,\n                         conjugate_a=conj)\n    x = triangular_solve(lu, x, left_side=True, lower=True, unit_diagonal=True,\n                         transpose_a=True, conjugate_a=conj)\n    _, ind = lax.sort_key_val(permutation, lax.iota('int32', permutation.shape[0]))\n    x = x[ind, :]\n  else:\n    raise ValueError(f\"'trans' value must be 0, 1, or 2, got {trans}\")\n  return lax.reshape(x, b.shape)\n\n\n@api.jit(static_argnums=(3,))\ndef _lu_solve(lu: Array, permutation: Array, b: Array, trans: int) -> Array:\n  if len(lu.shape) < 2 or lu.shape[-1] != lu.shape[-2]:\n    raise ValueError(\"last two dimensions of LU decomposition must be equal, \"\n                     \"got shape {}\".format(lu.shape))\n  if len(b.shape) < 1:\n    raise ValueError(\"b matrix must have rank >= 1, got shape {}\"\n                     .format(b.shape))\n  # Broadcasting follows NumPy's convention for linalg.solve: the RHS is\n  # treated as a (batched) vector if the number of dimensions differ by 1.\n  # Otherwise, broadcasting rules apply.\n  rhs_vector = lu.ndim == b.ndim + 1\n  if rhs_vector:\n    if b.shape[-1] != lu.shape[-1]:\n      raise ValueError(\"When LU decomposition matrix and b have the same \"\n                       \"number of dimensions, last axis of LU decomposition \"\n                       \"matrix (shape {}) and b array (shape {}) must match\"\n                       .format(lu.shape, b.shape))\n    b = b[..., np.newaxis]\n  else:\n    if b.shape[-2] != lu.shape[-1]:\n      raise ValueError(\"When LU decomposition matrix and b different \"","sourceCodeStart":1812,"sourceCodeEnd":1848,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/linalg.py#L1812-L1848","documentation":"jax/_src/lax/linalg.py:1830 in _lu_solve (public lu_solve). The precomputed LU factorization must be a square matrix per batch: shape [..., n, n] with ndim >= 2. Passing a vector, scalar, or rectangular array raises ValueError.","triggerScenarios":"Calling jax.lax.linalg.lu_solve with lu that is 0-d/1-d or whose last two dims differ — e.g. passing the original matrix A instead of the factorized lu output, or slicing the lu result incorrectly across batch dims.","commonSituations":"Mixing up arguments of (lu, permutation, b) tuples from jax.scipy.linalg.lu_factor / lu; caching a stale factorization of different shape after b changed; splitting batched factors along the wrong axis.","solutions":["Factor first: lu, piv = jax.scipy.linalg.lu_factor(a); then lu_solve(lu, piv, b)","Check lu.ndim >= 2 and lu.shape[-1] == lu.shape[-2] with an assert before solving","Re-factor whenever the matrix shape changes instead of reusing cached factors"],"exampleFix":"// before\nx = jax.lax.linalg.lu_solve(A, p, b)  # A raw matrix, not factors\n// after\nlu, piv = jax.scipy.linalg.lu_factor(A)\nx = jax.lax.linalg.lu_solve(lu, piv, b)","handlingStrategy":"validation","validationCode":"assert lu.ndim >= 2 and lu.shape[-1] == lu.shape[-2], lu.shape","typeGuard":"def is_lu_factor(lu: jax.Array) -> bool:\n    return lu.ndim >= 2 and lu.shape[-1] == lu.shape[-2]","tryCatchPattern":null,"preventionTips":["Only pass outputs of lu_factor to lu_solve","Re-factor when matrix shape changes"],"tags":["jax","lu-solve","shape-validation","square-matrix"],"backgroundTag":"matrix-not-square","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}