jax-ml/jax · error · TypeError

transpose_solve required for backwards mode automatic differ

Error message

transpose_solve required for backwards mode automatic differentiation of custom_linear_solve

What it means

jax.lax.custom_linear_solve differentiates through either (a) solving another system with the same matrix, or (b) explicitly calling a user-provided transpose_solve. The transpose rule requires jaxprs.transpose_solve; if custom_linear_solve was called without transpose_solve and matrix_symmetric=True was not applicable (so no transpose jaxpr was built), backward-mode differentiation raises TypeError.

Source

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

        core.jaxpr_as_fun(jaxprs.matvec), params.matvec, params_dot.matvec,
        jaxprs.matvec.debug_info, *x_leaves)
    rhs = _map(ad.add_tangents, b_dot, _map(operator.neg, matvec_tangents))

  x_dot = linear_solve_p.bind(*(_flatten(params) + rhs), **kwargs)

  # split into x tangents and aux tangents (these become zero)
  dx_leaves, daux_leaves = split_list(x_dot, [num_x_leaves])

  daux_leaves = _map(ad_util.p2tz, daux_leaves)

  x_dot = dx_leaves + daux_leaves

  return x, x_dot


def _linear_solve_transpose_rule(cotangent, *primals, const_lengths, jaxprs):
  if jaxprs.transpose_solve is None:
    raise TypeError('transpose_solve required for backwards mode automatic '
                    'differentiation of custom_linear_solve')

  params, b = _split_linear_solve_args(primals, const_lengths)
  if any(ad.is_undefined_primal(x) for xs in params for x in xs):
    raise NotImplementedError("open an issue at https://github.com/google/jax !!")
  assert all(ad.is_undefined_primal(x) for x in b)  # TODO(mattjj): why?
  x_cotangent, other_cotangents = split_list(cotangent, [len(b)])
  if any(type(ct) is not ad_util.Zero for ct in other_cotangents):
    raise NotImplementedError("open an issue at https://github.com/google/jax !!")
  del other_cotangents
  x_cotangent_ = _map(ad_util.instantiate, x_cotangent)
  cotangent_b_full = linear_solve_p.bind(
      *_flatten(params.transpose()), *x_cotangent_,
      const_lengths=const_lengths.transpose(), jaxprs=jaxprs.transpose())
  cotangent_b, _ = split_list(cotangent_b_full, [len(b)])
  return [None] * sum(const_lengths) + cotangent_b

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass transpose_solve to custom_linear_solve, e.g. transpose_solve=lambda Mt, v: solve_for(Mt.T if Mt is not None else None, v) if symmetric; otherwise a solver for the transposed system
2. Set matrix_symmetric=True if the operator is symmetric so the same solve serves as transpose_solve
3. Use jax.jacrev with a forward-over-reverse scheme or jvp if only first derivatives of the solve w.r.t. b are needed

Example fix

// before
x = jax.lax.custom_linear_solve(matvec, b, solve=direct_solve)
loss = jax.grad(f)(x)
// after
x = jax.lax.custom_linear_solve(matvec, b, solve=direct_solve,
                                 transpose_solve=lambda Mt, v: direct_solve(Mt, v))
loss = jax.grad(f)(x)

# ensure transpose path exists before grad
assert matrix_symmetric or transpose_solve is not None, \
    'grad requires transpose_solve or matrix_symmetric=True'

try:
    jax.grad(f)(x)
except TypeError as e:
    if 'transpose_solve required' in str(e):
        f2 = remake_with_transpose_solve(f)
        return jax.grad(f2)(x)
    raise

When it happens

Trigger: Calling jax.grad (or jax.vjp) on a function containing custom_linear_solve where neither transpose_solve was passed nor matrix_symmetric=True was set, and the solve function is not implicitly transposable.

Common situations: Wrapping external solvers (CuSPARSE, PETSc, scipy with differing transpose semantics) for use under grad; forward-mode (jvp) works but switching to loss.backward-style vjp exposes the missing transpose path.

Related errors

• Pure callbacks do not support transpose. Please use `jax.cus
• for transpose support, subclass {type(self)} must implement
• {type(_prim).__name__}.transpose should return None or a dic
• open an issue at https://github.com/google/jax !!
• transpose with implicit broadcasting of unshaped values. Got

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