{"record":{"id":"84324bd7ab708611","repo":"jax-ml/jax","slug":"can-t-use-defjvps-with-nondiff-argnums-84324b","errorCode":null,"errorMessage":"Can't use ``defjvps`` with ``nondiff_argnums``.","messagePattern":"Can't use ``defjvps`` with ``nondiff_argnums``\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":1183,"sourceCode":"\n\nclass custom_jvp3:\n  jvp_fun: Callable | None = None\n  symz: bool = False\n\n  def __init__(self, f, nondiff_argnums=(), nondiff_argnames=()):\n    self.static_argnums = _set_up_nondiff(f, nondiff_argnums, nondiff_argnames)\n    update_wrapper(self, f)\n    self.f = f\n\n  def defjvp(self, jvp, symbolic_zeros=False):\n    self.jvp_fun = jvp\n    self.symz = symbolic_zeros\n    return jvp\n\n  def defjvps(self, *jvps):\n    if self.static_argnums:\n      raise TypeError(\"Can't use ``defjvps`` with ``nondiff_argnums``.\")\n    def jvp(primals, tangents):\n      primal_out = self(*primals)\n      zeros = tree_map(ad_util.p2tz, primal_out)\n      all_tangents_out = [j(t, primal_out, *primals) if j else zeros\n                          for t, j in zip(tangents, jvps)]\n      sum_tangents = lambda _, x, *xs: reduce(ad.add_tangents, xs, x)\n      tangent_out = tree_map(sum_tangents, primal_out, *all_tangents_out)\n      return primal_out, tangent_out\n    self.defjvp(jvp)\n\n  def __call__(self, *args, **kwargs):\n    if not self.jvp_fun:\n      msg = (f\"No JVP defined for custom_jvp function {self.f.__name__} \"\n             \"using defjvp.\")\n      raise AttributeError(msg)\n\n    try:\n      args = resolve_kwargs(self.f, args, kwargs)","sourceCodeStart":1165,"sourceCodeEnd":1201,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L1165-L1201","documentation":"The defjvps convenience method (per-argument JVP definitions) only works when no nondiff_argnums were specified on the @jax.custom_jvp decorator, because defjvps builds a rule that differentiates with respect to every positional argument. Combining the two is a TypeError.","triggerScenarios":"Calling f.defjvps(jvp1, jvp2, ...) on a function decorated as @jax.custom_jvp(nondiff_argnums=(1,)).","commonSituations":"Porting code between the defjvp and defjvps styles; adding nondiff_argnums later for static config without removing defjvps usage.","solutions":["Use defjvp with a single full JVP function instead of defjvps","Remove nondiff_argnums from the decorator if per-argument definition is desired","Move static parameters outside the function (close over them or use functools.partial)"],"exampleFix":"// before\n@jax.custom_jvp(nondiff_argnums=(1,))\ndef f(x, n): ...\nf.defjvps(lambda t, p, x: 2*t)   # TypeError\n// after\n@jax.custom_jvp\ndef f(x, n): ...\n@f.defjvp\ndef f_jvp(primals, tangents):\n  (x, n), (xd, nd) = primals, tangents\n  return f(x, n), 2 * xd","handlingStrategy":"validation","validationCode":"if getattr(f, 'static_argnums', None):\n    raise TypeError('use defjvp, not defjvps, with nondiff_argnums')\nf.defjvps(*jvps)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Don't mix nondiff_argnums with defjvps","Move static params outside the decorated function when using defjvps"],"tags":["jax","custom-jvp","defjvps","api-misuse"],"backgroundTag":"jax-custom-jvp-defjvps-nondiff-conflict","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}