jax-ml/jax · error · NotImplementedError

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

Error message

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

What it means

HiPrim's default forward-mode (jvp) rule is a stub. Subclasses used under jax.jvp or jax.linearize must implement `jvp(primals, tangents)`.

Source

Thrown at jax/_src/hijax.py:190

  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` "
        "and `linearized`, or derive them from its `jvp` rule by setting "
        "`lin, linearized = linearize_from_jvp`")

  # optional transpose rule, for primitives that are linear in some inputs
  def transpose(self, out_ct, *maybe_accums):
    raise NotImplementedError(f"for transpose support, subclass {type(self)} "
                              "must implement `transpose`")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Implement `def jvp(self, primals, tangents)` on the subclass
2. Or derive forward rules from linearize via `lin, linearized = linearize_from_jvp` if only jvp exists (this error means jvp itself is missing, so implement it)
3. Avoid jvp/jacfwd/linearize over this primitive

Example fix

class MyPrim(hijax.HiPrim):
  # after
  def jvp(self, primals, tangents):
    out = self(*primals)
    return out, tree_map(jnp.zeros_like, out)  # placeholder

if type(prim).jvp is hijax.HiPrim.jvp:
    raise ValueError(f'{type(prim).__name__} lacks a jvp rule')

def has_jvp_rule(p) -> bool:
    return type(p).jvp is not hijax.HiPrim.jvp

try:
    jax.jvp(f, (x,), (t,))
except NotImplementedError as e:
    if 'jvp' in str(e):  # fall back to reverse mode
        return jax.grad(f)(x)
    raise

When it happens

Trigger: Calling jax.jvp (or jax.linearize/jacfwd) on a function that applies a HiPrim subclass without a jvp method.

Common situations: Writing a hijax custom primitive with only a primal implementation and then differentiating it in forward mode, or using scan/ode solvers that internally use jvp.

Related errors

• for grad support, subclass {type(self)} must implement `vjp_
• for linearize support, subclass {type(self)} must implement
• for transpose support, subclass {type(self)} must implement
• open an issue at https://github.com/google/jax !!
• Unimplemented group_offset support.

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