jax-ml/jax · error · ValueError

When LU decomposition matrix and b different numbers of dime

Error message

When LU decomposition matrix and b different numbers of dimensions, last axis of LU decomposition matrix (shape {lu.shape}) and second to last axis of b array (shape {b.shape}) must match

What it means

jax/_src/lax/linalg.py:1848 in _lu_solve (public lu_solve). When b is a (batched) matrix with the same ndim handling as lu (matrix RHS branch), its second-to-last axis must match the LU size: b.shape[-2] == lu.shape[-1] (i.e. b is [..., n, k]). Any other width raises this ValueError.

Source

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

    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:
    fn = api.vmap(fn, in_axes=(0, 0, 0, None))
  x = fn(lu, permutation, b, trans)
  return x[..., 0] if rhs_vector else x

# Support operation for LU decomposition: Transformation of the pivots returned
# by LU decomposition into permutations.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Transpose b so the solve dimension is -2: b = b.mT (or b.T for 2-d)
2. Re-factor and rebuild b together so shapes stay consistent
3. Assert b.shape[-2] == lu.shape[-1] before solving

Example fix

// before
x = jax.lax.linalg.lu_solve(lu, piv, b)  # b: (k, n) but needs (n, k)
// after
x = jax.lax.linalg.lu_solve(lu, piv, b.mT)

assert b.shape[-2] == lu.shape[-1], (lu.shape, b.shape)
# fix convention: b = b.mT if transposed

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

When it happens

Trigger: Solving with multiple right-hand sides where b is (..., k, n) transposed instead of (..., n, k), or b built from a different-sized matrix; also mismatched batch broadcasting after reshaping.

Common situations: Transposed RHS convention differences when porting scipy.linalg.lu_solve / numpy linalg.solve code; batching where b was concatenated along the wrong axis; using a stale factorization with a new wider b.

Related errors

• When LU decomposition matrix and b have the same number of d
• 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/de8a08490490cb65. Report an issue: GitHub.