jax-ml/jax · error · ValueError

last two dimensions of LU decomposition must be equal, got s

Error message

last two dimensions of LU decomposition must be equal, got shape {lu.shape}

What it means

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.

Source

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

    x = triangular_solve(lu, x, left_side=True, lower=True, unit_diagonal=True)
    x = triangular_solve(lu, x, left_side=True, lower=False)
  elif trans == 1 or trans == 2:
    conj = trans == 2
    x = triangular_solve(lu, x, left_side=True, lower=False, transpose_a=True,
                         conjugate_a=conj)
    x = triangular_solve(lu, x, left_side=True, lower=True, unit_diagonal=True,
                         transpose_a=True, conjugate_a=conj)
    _, ind = lax.sort_key_val(permutation, lax.iota('int32', permutation.shape[0]))
    x = x[ind, :]
  else:
    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 "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Factor first: lu, piv = jax.scipy.linalg.lu_factor(a); then lu_solve(lu, piv, b)
2. Check lu.ndim >= 2 and lu.shape[-1] == lu.shape[-2] with an assert before solving
3. Re-factor whenever the matrix shape changes instead of reusing cached factors

Example fix

// before
x = jax.lax.linalg.lu_solve(A, p, b)  # A raw matrix, not factors
// after
lu, piv = jax.scipy.linalg.lu_factor(A)
x = jax.lax.linalg.lu_solve(lu, piv, b)

assert lu.ndim >= 2 and lu.shape[-1] == lu.shape[-2], lu.shape

def is_lu_factor(lu: jax.Array) -> bool:
    return lu.ndim >= 2 and lu.shape[-1] == lu.shape[-2]

When it happens

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

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

Related errors

• Argument to symmetric eigendecomposition must have shape [..
• b matrix must have rank >= 1, got shape {b.shape}
• When LU decomposition matrix and b have the same number of d
• When LU decomposition matrix and b different numbers of dime
• scan got `length` argument of {} which disagrees with leadin

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