{"record":{"id":"40a84f4fd5d6b95a","repo":"jax-ml/jax","slug":"for-jvp-support-subclass-type-self-must-implem","errorCode":null,"errorMessage":"for jvp support, subclass {type(self)} must implement `jvp`","messagePattern":"for jvp support, subclass (.+?) must implement `jvp`","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":190,"sourceCode":"  def vjp_bwd(self, res, outgrad, /, *arg_accums):\n    if self.vjp_bwd_retval_logs:\n      args_grad, logs = self.vjp_bwd_retval(res, outgrad)\n    else:\n      args_grad, logs = self.vjp_bwd_retval(res, outgrad), None\n    maybe_accum = lambda acc, v: isinstance(acc, ad.GradAccum) and acc.accum(v)\n    tree_map(maybe_accum, arg_accums, args_grad)\n    return logs\n\n  def vjp_bwd_retval(self, res, outgrad, /):\n    # Classic API: returns values instead of using accumulators\n    raise NotImplementedError(\n        f\"for grad support, subclass {type(self)} must implement `vjp_bwd` or \"\n        \"`vjp_bwd_retval`, or derive its reverse-mode rules by setting \"\n        \"`vjp_fwd, vjp_bwd_retval = vjp_from_jvp` (or `= vjp_from_lin`)\")\n\n  # optional forward-mode AD interfaces\n  def jvp(self, primals, tangents):\n    raise NotImplementedError(f\"for jvp support, subclass {type(self)} must \"\n                              \"implement `jvp`\")\n\n  def lin(self, nzs_in, *primals):\n    raise NotImplementedError(\n        f\"for linearize support, subclass {type(self)} must implement `lin` \"\n        \"and `linearized`, or derive them from its `jvp` rule by setting \"\n        \"`lin, linearized = linearize_from_jvp`\")\n\n  def linearized(self, residuals, *tangents):\n    raise NotImplementedError(\n        f\"for linearize support, subclass {type(self)} must implement `lin` \"\n        \"and `linearized`, or derive them from its `jvp` rule by setting \"\n        \"`lin, linearized = linearize_from_jvp`\")\n\n  # optional transpose rule, for primitives that are linear in some inputs\n  def transpose(self, out_ct, *maybe_accums):\n    raise NotImplementedError(f\"for transpose support, subclass {type(self)} \"\n                              \"must implement `transpose`\")","sourceCodeStart":172,"sourceCodeEnd":208,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L172-L208","documentation":"HiPrim's default forward-mode (jvp) rule is a stub. Subclasses used under jax.jvp or jax.linearize must implement `jvp(primals, tangents)`.","triggerScenarios":"Calling jax.jvp (or jax.linearize/jacfwd) on a function that applies a HiPrim subclass without a jvp method.","commonSituations":"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.","solutions":["Implement `def jvp(self, primals, tangents)` on the subclass","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)","Avoid jvp/jacfwd/linearize over this primitive"],"exampleFix":"class MyPrim(hijax.HiPrim):\n  # after\n  def jvp(self, primals, tangents):\n    out = self(*primals)\n    return out, tree_map(jnp.zeros_like, out)  # placeholder\n","handlingStrategy":"validation","validationCode":"if type(prim).jvp is hijax.HiPrim.jvp:\n    raise ValueError(f'{type(prim).__name__} lacks a jvp rule')","typeGuard":"def has_jvp_rule(p) -> bool:\n    return type(p).jvp is not hijax.HiPrim.jvp","tryCatchPattern":"try:\n    jax.jvp(f, (x,), (t,))\nexcept NotImplementedError as e:\n    if 'jvp' in str(e):  # fall back to reverse mode\n        return jax.grad(f)(x)\n    raise","preventionTips":["Provide jvp for any primitive that may be used with jacfwd/linearize","Test primitives under both grad and jvp"],"tags":["jax","autodiff","forward-mode","not-implemented"],"backgroundTag":"missing-autodiff-rule","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}