jax-ml/jax · error · ValueError

{}() output shapes must match {}, got {} and {}

Error message

{}() output shapes must match {}, got {} and {}

What it means

custom_linear_solve (and custom_root) validate that the user-supplied solve function returns arrays whose shapes exactly match the shape of the expected solution (e.g. the b matrix). The helper _check_shapes compares actual vs expected shape lists and raises ValueError showing both when they differ.

Source

Thrown at jax/_src/lax/control_flow/solves.py:205

def _transpose_one_output(linear_fun, primals):
  transpose_fun = api.linear_transpose(linear_fun, primals)
  def transposed_fun(x):
    (y,) = transpose_fun(x)
    return y
  return transposed_fun


def _flatten(args):
  return [x for arg in args for x in arg]


def _check_shapes(func_name, expected_name, actual, expected):
  actual_shapes = _map(np.shape, actual)
  expected_shapes = _map(np.shape, expected)
  if actual_shapes != expected_shapes:
    raise ValueError(
        f"{func_name}() output shapes must match {expected_name}, "
        f"got {actual_shapes} and {expected_shapes}")


@partial(api_boundary, repro_api_name="jax.custom_linear_solve")
def custom_linear_solve(
    matvec: Callable,
    b: Any,
    solve: Callable[[Callable, Any], Any],
    transpose_solve: Callable[[Callable, Any], Any] | None = None,
    symmetric=False, has_aux=False):
  """Perform a matrix-free linear solve with implicitly defined gradients.

  This function allows for overriding or defining gradients for a linear
  solve directly via implicit differentiation at the solution, rather than by
  differentiating *through* the solve operation. This can sometimes be much faster
  or more numerically stable, or differentiating through the solve operation
  may not even be implemented (e.g., if ``solve`` uses ``lax.while_loop``).

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the solve function return exactly the same shapes (and pytree structure) as b
2. Add assert statements inside solve: lambda A, b: (assert matching shapes, x)
3. Return x.reshape(b.shape) as a defensive last line of the solve closure

Example fix

// before
def solve(A, b):
  return jsp.linalg.solve(A, b).squeeze()  # wrong shape for b of shape (n, 1)
// after
def solve(A, b):
  return jsp.linalg.solve(A, b).reshape(b.shape)

x = solve(A, b)
actual = [np.shape(a) for a in jax.tree_util.tree_leaves(x)]
expected = [np.shape(e) for e in jax.tree_util.tree_leaves(b)]
assert actual == expected, (actual, expected)

When it happens

Trigger: Passing a `solve`/`solve_transpose` function to jax.lax.custom_linear_solve that returns a reshaped or transposed result, or the wrong number of arrays, e.g. returning b.T shape or a scalar instead of a vector.

Common situations: Custom iterative/Cholesky solvers returning solutions with extra/missing batch dims; a matvec closure that accidentally squeezes the batch dimension; mismatch between the number of outputs and the number of b arrays in pytree input.

Related errors

• unexpected JAX type (e.g. shape/dtype) for argument to VJP f
• cotangent type does not match function output, expected {out
• Mismatched number of outputs from callback. Expected: {}, Ac
• Incorrect output shape for return value #{i}: Expected: {out
• Length of sharding.spec ({len(out_s.spec)}) must be equal to

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