jax-ml/jax · error · ValueError

When LU decomposition matrix and b have the same number of d

Error message

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

What it means

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.

Source

Thrown at jax/_src/lax/linalg.py:1841

    raise ValueError(f"'trans' value must be 0, 1, or 2, got {trans}")
  return lax.reshape(x, b.shape)


@api.jit(static_argnums=(3,))
def _lu_solve(lu: Array, permutation: Array, b: Array, trans: int) -> Array:
  if len(lu.shape) < 2 or lu.shape[-1] != lu.shape[-2]:
    raise ValueError("last two dimensions of LU decomposition must be equal, "
                     "got shape {}".format(lu.shape))
  if len(b.shape) < 1:
    raise ValueError("b matrix must have rank >= 1, got shape {}"
                     .format(b.shape))
  # Broadcasting follows NumPy's convention for linalg.solve: the RHS is
  # treated as a (batched) vector if the number of dimensions differ by 1.
  # Otherwise, broadcasting rules apply.
  rhs_vector = lu.ndim == b.ndim + 1
  if rhs_vector:
    if b.shape[-1] != lu.shape[-1]:
      raise ValueError("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"
                       .format(lu.shape, b.shape))
    b = b[..., np.newaxis]
  else:
    if b.shape[-2] != lu.shape[-1]:
      raise ValueError("When LU decomposition matrix and b different "
                       "numbers of dimensions, last axis of LU decomposition "
                       "matrix (shape {}) and second to last axis of b array "
                       "(shape {}) must match"
                       .format(lu.shape, b.shape))

  batch_shape = lax.broadcast_shapes(lu.shape[:-2], permutation.shape[:-1], b.shape[:-2])
  lu = _broadcast_to(lu, (*batch_shape, *lu.shape[-2:]))
  permutation = _broadcast_to(permutation, (*batch_shape, permutation.shape[-1]))
  b = _broadcast_to(b, (*batch_shape, *b.shape[-2:]))
  fn = _lu_solve_core
  for _ in batch_shape:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Verify b.shape[-1] == lu.shape[-1]; construct b from the same matrix's rows/columns
2. If b is a matrix RHS, ensure it keeps 2 dims so the matrix branch is used
3. Pad or slice b to length n if the mismatch is a known padding artifact

Example fix

// before
x = jax.lax.linalg.lu_solve(lu, piv, b)  # b: (m,), lu: (n, n), m != n
// after
assert b.shape[-1] == lu.shape[-1]
x = jax.lax.linalg.lu_solve(lu, piv, b)

assert b.shape[-1] == lu.shape[-1], (lu.shape, b.shape)

def lu_vector_rhs_ok(lu, b) -> bool:
    return b.ndim == lu.ndim - 1 and b.shape[-1] == lu.shape[-1]

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• When LU decomposition matrix and b different numbers of dime
• last two dimensions of LU decomposition must be equal, got s
• b matrix must have rank >= 1, got shape {b.shape}
• scan got `length` argument of {} which disagrees with leadin
• conv_general_dilated batch_group_count must divide lhs batch

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/d4a06c70e9d4ac66. Report an issue: GitHub.