jax-ml/jax · error · ValueError

'trans' value must be 0, 1, or 2, got {trans}

Error message

'trans' value must be 0, 1, or 2, got {trans}

What it means

jax/_src/lax/linalg.py:1823 in _lu_solve_core. The trans parameter selects the solve variant: 0 = A x = b, 1 = A^T x = b, 2 = A^H x = b. Any other integer falls through the if/elif chain and raises ValueError — it is a static Python int (the function is jitted with static_argnums for it), so this is a plain caller bug.

Source

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

def _lu_solve_core(lu: Array, permutation: Array, b: Array, trans: int) -> Array:
  m = lu.shape[0]
  x = lax.reshape(b, (m, math.prod(b.shape[1:])))
  if trans == 0:
    x = x[permutation, :]
    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 "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Map your convention explicitly: {'N': 0, 'T': 1, 'C': 2}[mode] before calling
2. Validate trans at the API boundary: assert trans in (0, 1, 2)

Example fix

// before
x = jax.lax.linalg.lu_solve(lu, p, b, trans='T')
// after
x = jax.lax.linalg.lu_solve(lu, p, b, trans={'N':0,'T':1,'C':2}['T'])

TRANS = {'N': 0, 'T': 1, 'C': 2}
assert trans in (0, 1, 2)

def valid_trans(t) -> bool:
    return t in (0, 1, 2)

When it happens

Trigger: Calling jax.lax.linalg.lu_solve(lu, permutation, b, trans) with trans not in {0, 1, 2}, e.g. passing -1, 3, a bool, or a string like 'T' (SciPy convention).

Common situations: Porting SciPy/BLAS conventions where trans is 'N'/'T'/'C' or 0-indexed enums differ; passing a NumPy integer or config flag that drifts from expected values after a refactor.

Related errors

• size and fill_value arguments cannot be used in three-term w
• {name} wrapped function must be passed at least one argument
• primal and tangent arguments to jax.jvp must be tuples or li
• check_error takes an Error as argument, got type {type(error
• Can't use ``defjvps`` with ``nondiff_argnums``.

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