jax-ml/jax · error · NotImplementedError

open an issue at https://github.com/google/jax !!

Error message

open an issue at https://github.com/google/jax !!

What it means

Inside the custom_linear_solve transpose rule, after splitting primals into params and b, the rule asserts that only the b (right-hand-side) arguments carry tangent/undefined-primal status. If any matrix/operator parameter itself is being differentiated through (an 'open' AD input in the transpose), the rule hits an unimplemented path and raises NotImplementedError asking for a GitHub issue.

Source

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

  # 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


def _linear_solve_batching_rule(axis_data, args, dims, const_lengths, jaxprs):
  orig_bat = [d is not None for d in dims]

  params, b = _split_linear_solve_args(args, const_lengths)
  params_dims, b_dims = _split_linear_solve_args(dims, const_lengths)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Provide an explicit transpose_solve that handles the operator's parameters so the generic path is bypassed
2. Stop-gradient the operator parameters: jax.lax.stop_gradient on the matvec's captured weights if you only need grads w.r.t. b or upstream inputs
3. Reparameterize so the operator is constant: precompute the matrix outside the traced function
4. If it looks like a genuine unsupported case, open the requested issue at github.com/google/jax

Example fix

// before
x = jax.lax.custom_linear_solve(lambda v: apply_A(A, v), b, solve)
jax.grad(f)(...)  # A differentiable -> transpose rule fails
// after
A_ = jax.lax.stop_gradient(A)
x = jax.lax.custom_linear_solve(lambda v: apply_A(A_, v), b, solve)

# stop-gradient operator params before the solve
matvec_safe = lambda v: matvec(jax.lax.stop_gradient(A), v)

When it happens

Trigger: Taking jax.grad of a function where custom_linear_solve's matvec closure captures differentiable (traced) arrays — i.e. differentiating w.r.t. the linear operator's parameters, in a configuration not covered by the implicit-diff path (e.g. no transpose_solve dealing with params).

Common situations: End-to-end differentiation through a neural PDE solver or learned preconditioner used inside custom_linear_solve without providing transpose_solve; jvp works but vjp/grad hits this.

Related errors

• for grad support, subclass {type(self)} must implement `vjp_
• for jvp support, subclass {type(self)} must implement `jvp`
• for linearize support, subclass {type(self)} must implement
• for transpose support, subclass {type(self)} must implement
• transpose_solve required for backwards mode automatic differ

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