{"record":{"id":"18aba22705aad887","repo":"jax-ml/jax","slug":"custom-jvp-inputs-marked-with-nondiff-argnums-must","errorCode":null,"errorMessage":"custom_jvp inputs marked with nondiff_argnums must be static, not Tracers","messagePattern":"custom_jvp inputs marked with nondiff_argnums must be static, not Tracers","errorType":"exception","errorClass":"UnexpectedTracerError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":1208,"sourceCode":"      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)\n    except TypeError as e:\n      raise TypeError(\n          \"The input arguments to the custom_jvp-decorated function \"\n          f\"{self.f.__name__} could not be resolved to positional-only \"\n          f\"arguments. Binding failed with the error:\\n{e}\") from e\n    if any(isinstance(args[i], core.Tracer) for i in self.static_argnums):\n      raise UnexpectedTracerError(\"custom_jvp inputs marked with nondiff_argnums \"\n                                  \"must be static, not Tracers\")\n    if all(is_hashable(args[i]) for i in self.static_argnums):\n      traced = api.jit(self.f, static_argnums=(*self.static_argnums,)).trace(*args)\n    else:\n      # jit requires hashable static_argnums values, but classic custom_jvp\n      # accepted unhashable nondiff_argnums values, so close over them instead\n      which_static = [i in self.static_argnums for i in range(len(args))]\n      dyn_args, static_args = partition_list(which_static, args)\n      f = dyn_args_fun(self.f, self.static_argnums,\n                       tuple(map(WrapHashably, static_args)), len(args))\n      traced = api.jit(f).trace(*dyn_args)\n    if any(isinstance(x, core.Tracer) for x in traced._consts):\n      t = next(x for x in traced._consts if isinstance(x, core.Tracer))\n      raise UnexpectedTracerError(\n          f\"custom_jvp-decorated function {self.f} closed over a {type(t).__name__} \"\n          f\"of type {t.aval.str_short()}, but custom_jvp functions can't close \"\n          f\"over Tracers. Rewrite {self.f} to take it as an explicit input.\")\n    args = tuple(Static(x) if i in self.static_argnums else x for i, x in enumerate(args))","sourceCodeStart":1190,"sourceCodeEnd":1226,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L1190-L1226","documentation":"custom_jvp analogue of error 644: an argument listed in nondiff_argnums was a core.Tracer (produced by jit/grad/vmap/scan) at call time, but nondiff args must be static Python values. JAX raises UnexpectedTracerError before tracing.","triggerScenarios":"Calling a @jax.custom_jvp(nondiff_argnums=...) function inside jit/grad with a traced array in one of the static positions, e.g. passing a computed index array instead of a Python int.","commonSituations":"Same as 644: passing array-derived metadata (sizes, flags) that became tracers; loops where static args are accidentally results of traced ops; switching a previously-static config to a traced value.","solutions":["Convert to a concrete Python value before the call (int(n), .item() outside jit)","Move the argument into the differentiable args and handle it inside the rule","Hoist computation of the static value outside the transforming function"],"exampleFix":"// before\n@jax.custom_jvp(nondiff_argnums=(1,))\ndef f(x, n): ...\njit(lambda x, n: f(x, n))(x, traced_n)\n// after\nn_static = int(traced_n)  # computed outside jit\njit(lambda x: f(x, n_static))(x)","handlingStrategy":"type-guard","validationCode":"from jax.core import Tracer\nassert not any(isinstance(args[i], Tracer) for i in f.static_argnums)","typeGuard":"from jax.core import Tracer\ndef is_static(v) -> bool:\n    return not isinstance(v, Tracer)","tryCatchPattern":null,"preventionTips":["Convert traced scalars to Python ints outside jit","Keep nondiff args concrete Python values"],"tags":["jax","custom-jvp","tracer-error","nondiff-argnums"],"backgroundTag":"jax-unexpected-tracer-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}