{"record":{"id":"b83edb291261033d","repo":"jax-ml/jax","slug":"for-grad-support-subclass-type-self-must-imple-b83edb","errorCode":null,"errorMessage":"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`)","messagePattern":"for grad support, subclass (.+?) 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`\\)","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":183,"sourceCode":"        f\"for grad support, subclass {type(self)} must implement `vjp_fwd`, \"\n        \"or derive its reverse-mode rules from its jvp or lin rules by \"\n        \"setting `vjp_fwd, vjp_bwd_retval = vjp_from_jvp` (or `= \"\n        \"vjp_from_lin`)\")\n\n  vjp_bwd_retval_logs: bool = False\n\n  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` \"","sourceCodeStart":165,"sourceCodeEnd":201,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L165-L201","documentation":"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`).","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"class MyPrim(hijax.HiPrim):\n  # before: no vjp rules\n# after\n  def vjp_fwd(self, *args):\n    out = self(*args)\n    return out, None\n  def vjp_bwd_retval(self, res, out_ct):\n    return ...  # input cotangents","handlingStrategy":"validation","validationCode":"missing = (type(prim).vjp_bwd is hijax.HiPrim.vjp_bwd and type(prim).vjp_bwd_retval is hijax.HiPrim.vjp_bwd_retval)\nif missing and needs_grad:\n    raise ValueError(f'{type(prim).__name__} lacks VJP rules; cannot differentiate')","typeGuard":"def has_vjp_rules(p) -> bool:\n    return not (type(p).vjp_bwd is hijax.HiPrim.vjp_bwd and\n                type(p).vjp_bwd_retval is hijax.HiPrim.vjp_bwd_retval)","tryCatchPattern":"try:\n    jax.grad(f)(x)\nexcept NotImplementedError as e:\n    if 'vjp_bwd' in str(e):\n        raise RuntimeError('primitive lacks backward rule: ' + str(e)) from e\n    raise","preventionTips":["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"],"tags":["jax","autodiff","custom-primitive","not-implemented"],"backgroundTag":"missing-autodiff-rule","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}