{"record":{"id":"d4a06c70e9d4ac66","repo":"jax-ml/jax","slug":"when-lu-decomposition-matrix-and-b-have-the-same-n","errorCode":null,"errorMessage":"When LU decomposition matrix and b have the same number of dimensions, last axis of LU decomposition matrix (shape {lu.shape}) and b array (shape {b.shape}) must match","messagePattern":"When LU decomposition matrix and b have the same number of dimensions, last axis of LU decomposition matrix \\(shape (.+?)\\) and b array \\(shape (.+?)\\) must match","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/linalg.py","lineNumber":1841,"sourceCode":"    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 \"\n                       \"numbers of dimensions, last axis of LU decomposition \"\n                       \"matrix (shape {}) and second to last axis of b array \"\n                       \"(shape {}) must match\"\n                       .format(lu.shape, b.shape))\n\n  batch_shape = lax.broadcast_shapes(lu.shape[:-2], permutation.shape[:-1], b.shape[:-2])\n  lu = _broadcast_to(lu, (*batch_shape, *lu.shape[-2:]))\n  permutation = _broadcast_to(permutation, (*batch_shape, permutation.shape[-1]))\n  b = _broadcast_to(b, (*batch_shape, *b.shape[-2:]))\n  fn = _lu_solve_core\n  for _ in batch_shape:","sourceCodeStart":1823,"sourceCodeEnd":1859,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/linalg.py#L1823-L1859","documentation":"jax/_src/lax/linalg.py:1841 in _lu_solve (public lu_solve). When b is treated as a batched vector (lu.ndim == b.ndim + 1), the vector length must equal the matrix size: b.shape[-1] == lu.shape[-1]. Mismatch raises this ValueError, mirroring NumPy linalg.solve's vector convention.","triggerScenarios":"Solving A x = b_vec where len(b_vec) != n, e.g. lu from an (n, n) matrix but b of length m != n; also when b was meant to be a (n, k) matrix but got squeezed to the wrong length.","commonSituations":"Residual/normal-equation pipelines where b is computed from a differently-shaped matrix; transposing bugs (row vs column vector of wrong length); batched problems where one item's b has a different size after ragged padding.","solutions":["Verify b.shape[-1] == lu.shape[-1]; construct b from the same matrix's rows/columns","If b is a matrix RHS, ensure it keeps 2 dims so the matrix branch is used","Pad or slice b to length n if the mismatch is a known padding artifact"],"exampleFix":"// before\nx = jax.lax.linalg.lu_solve(lu, piv, b)  # b: (m,), lu: (n, n), m != n\n// after\nassert b.shape[-1] == lu.shape[-1]\nx = jax.lax.linalg.lu_solve(lu, piv, b)","handlingStrategy":"validation","validationCode":"assert b.shape[-1] == lu.shape[-1], (lu.shape, b.shape)","typeGuard":"def lu_vector_rhs_ok(lu, b) -> bool:\n    return b.ndim == lu.ndim - 1 and b.shape[-1] == lu.shape[-1]","tryCatchPattern":null,"preventionTips":["Build b from the same matrix dimension as the factorization"],"tags":["jax","lu-solve","shape-validation","broadcasting"],"backgroundTag":"invalid-matrix-dimensions","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}