jax-ml/jax · error · NotImplementedError

for grad support, subclass {type(self)} must implement `vjp_

Error message

for grad support, subclass {type(self)} must implement `vjp_bwd` or `vjp_bwd_retval`, or derive its reverse-mode rules by setting `vjp_fwd, vjp_bwd_retval = vjp_from_jvp` (or `= vjp_from_lin`)

What it means

HiPrim's default reverse-mode rule is unimplemented. The hijax custom-primitive base class requires any subclass used under grad/vjp to either implement vjp_bwd or vjp_bwd_retval, or derive both from a forward-mode rule via `vjp_fwd, vjp_bwd_retval = vjp_from_jvp` (or `vjp_from_lin`).

Source

Thrown at jax/_src/hijax.py:183

        f"for grad support, subclass {type(self)} must implement `vjp_fwd`, "
        "or derive its reverse-mode rules from its jvp or lin rules by "
        "setting `vjp_fwd, vjp_bwd_retval = vjp_from_jvp` (or `= "
        "vjp_from_lin`)")

  vjp_bwd_retval_logs: bool = False

  def vjp_bwd(self, res, outgrad, /, *arg_accums):
    if self.vjp_bwd_retval_logs:
      args_grad, logs = self.vjp_bwd_retval(res, outgrad)
    else:
      args_grad, logs = self.vjp_bwd_retval(res, outgrad), None
    maybe_accum = lambda acc, v: isinstance(acc, ad.GradAccum) and acc.accum(v)
    tree_map(maybe_accum, arg_accums, args_grad)
    return logs

  def vjp_bwd_retval(self, res, outgrad, /):
    # Classic API: returns values instead of using accumulators
    raise NotImplementedError(
        f"for grad support, subclass {type(self)} must implement `vjp_bwd` or "
        "`vjp_bwd_retval`, or derive its reverse-mode rules by setting "
        "`vjp_fwd, vjp_bwd_retval = vjp_from_jvp` (or `= vjp_from_lin`)")

  # optional forward-mode AD interfaces
  def jvp(self, primals, tangents):
    raise NotImplementedError(f"for jvp support, subclass {type(self)} must "
                              "implement `jvp`")

  def lin(self, nzs_in, *primals):
    raise NotImplementedError(
        f"for linearize support, subclass {type(self)} must implement `lin` "
        "and `linearized`, or derive them from its `jvp` rule by setting "
        "`lin, linearized = linearize_from_jvp`")

  def linearized(self, residuals, *tangents):
    raise NotImplementedError(
        f"for linearize support, subclass {type(self)} must implement `lin` "

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Implement vjp_bwd(res, outgrad, *arg_accums) or the classic vjp_bwd_retval(self, res, outgrad) on the subclass
Alternatively set `vjp_fwd, vjp_bwd_retval = vjp_from_jvp` in the class body if a jvp rule exists
Alternatively set `vjp_fwd, vjp_bwd_retval = vjp_from_lin` if linearize rules exist
If gradients are not needed, avoid grad/vjp over this primitive

Example fix

class MyPrim(hijax.HiPrim):
  # before: no vjp rules
# after
  def vjp_fwd(self, *args):
    out = self(*args)
    return out, None
  def vjp_bwd_retval(self, res, out_ct):
    return ...  # input cotangents

Defensive patterns

Strategy: validation

Validate before calling

missing = (type(prim).vjp_bwd is hijax.HiPrim.vjp_bwd and type(prim).vjp_bwd_retval is hijax.HiPrim.vjp_bwd_retval)
if missing and needs_grad:
    raise ValueError(f'{type(prim).__name__} lacks VJP rules; cannot differentiate')

Type guard

def has_vjp_rules(p) -> bool:
    return not (type(p).vjp_bwd is hijax.HiPrim.vjp_bwd and
                type(p).vjp_bwd_retval is hijax.HiPrim.vjp_bwd_retval)

Try / catch

try:
    jax.grad(f)(x)
except NotImplementedError as e:
    if 'vjp_bwd' in str(e):
        raise RuntimeError('primitive lacks backward rule: ' + str(e)) from e
    raise

Prevention

Implement VJP rules (or vjp_from_jvp derivation) in every HiPrim subclass as a convention
Add a unit test that runs jax.grad through each custom primitive

When it happens

Trigger: Calling jax.grad, jax.vjp, or any reverse-mode AD transform on a function that applies a HiPrim subclass that did not define any VJP rules.

Common situations: Authoring a custom primitive in jax.experimental.hijax (e.g. wrapping a kernel or solver) and forgetting the backward rule, then running gradient-based optimization on it.

Related errors

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