{"record":{"id":"852bbe3f49a1de71","repo":"jax-ml/jax","slug":"can-t-use-defjvps-with-nondiff-argnums","errorCode":null,"errorMessage":"Can't use ``defjvps`` with ``nondiff_argnums``.","messagePattern":"Can't use ``defjvps`` with ``nondiff_argnums``\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_derivatives.py","lineNumber":243,"sourceCode":"      None.\n\n    Examples:\n\n      >>> @jax.custom_jvp\n      ... def f(x, y):\n      ...   return jnp.sin(x) * y\n      ...\n      >>> f.defjvps(lambda x_dot, primal_out, x, y: jnp.cos(x) * x_dot * y,\n      ...           lambda y_dot, primal_out, x, y: jnp.sin(x) * y_dot)\n\n      >>> x = jnp.float32(1.0)\n      >>> y = jnp.float32(2.0)\n      >>> with jnp.printoptions(precision=2):\n      ...   print(jax.value_and_grad(f)(x, y))\n      (Array(1.68, dtype=float32), Array(1.08, dtype=float32))\n    \"\"\"\n    if self.nondiff_argnums:\n      raise TypeError(\"Can't use ``defjvps`` with ``nondiff_argnums``.\")\n\n    def jvp(primals, tangents):\n      primal_out = self(*primals)\n      zeros = _zeros_like_pytree(primal_out)\n      all_tangents_out = [jvp(t, primal_out, *primals) if jvp else zeros\n                          for t, jvp in zip(tangents, jvps)]\n      tangent_out = tree_map(_sum_tangents, primal_out, *all_tangents_out)\n      return primal_out, tangent_out\n\n    self.defjvp(jvp)\n\n  @partial(traceback_util.api_boundary,\n           repro_api_name=\"jax.custom_jvp.__call__\")\n  def __call__(self, *args: Any, **kwargs: Any) -> ReturnValue:\n    debug = debug_info(\"custom_jvp fun\", self.fun, args, kwargs,\n                       static_argnums=self.nondiff_argnums)\n    primal_name = debug.func_name\n    if not self.jvp:","sourceCodeStart":225,"sourceCodeEnd":261,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_derivatives.py#L225-L261","documentation":"jax.custom_jvp's defjvps convenience helper cannot be combined with nondiff_argnums, because defjvps assumes every positional argument is differentiable and gets a jvp entry. Supplying both is rejected with a TypeError.","triggerScenarios":"Declaring @jax.custom_jvp(nondiff_argnums=(1,)) on a function and then calling f.defjvps(...) to register per-argument JVPs.","commonSituations":"Porting code that used static integer flags via nondiff_argnums while keeping a defjvps-style rule table; growing an existing custom_jvp function with extra static params.","solutions":["Switch to the explicit defjvp form (f.defjvp) which handles nondiff args by receiving them as static values in the rule","Drop nondiff_argnums and pass the static values as arrays/traced values if semantically acceptable"],"exampleFix":"# before\n@jax.custom_jvp(nondiff_argnums=(1,))\ndef f(x, n): ...\nf.defjvps(lambda t, p, x, n: ...)\n# after\n@jax.custom_jvp\ndef f(x, n): ...\n@f.defjvp\ndef f_jvp(primals, tangents):\n    x, n = primals\n    tx, _ = tangents\n    return f(x, n), tx * dfdx(x, n)","handlingStrategy":"validation","validationCode":"# at decoration time choose one style:\n# either @jax.custom_jvp (no nondiff_argnums) + defjvps, or nondiff_argnums + defjvp","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Prefer explicit defjvp when static args are needed","Document the choice near the decorator"],"tags":["jax","custom-jvp","autodiff","api-misuse"],"backgroundTag":"incompatible-api-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}