{"record":{"id":"da1b1c53c851b444","repo":"jax-ml/jax","slug":"custom-jvp-decorated-function-self-f-closed-over","errorCode":null,"errorMessage":"custom_jvp-decorated function {self.f} closed over a {type(t).__name__} of type {t.aval.str_short()}, but custom_jvp functions can't close over Tracers. Rewrite {self.f} to take it as an explicit input.","messagePattern":"custom_jvp-decorated function (.+?) closed over a (.+?) of type (.+?), but custom_jvp functions can't close over Tracers\\. Rewrite (.+?) to take it as an explicit input\\.","errorType":"exception","errorClass":"UnexpectedTracerError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":1222,"sourceCode":"          \"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))\n    in_avals = tree_map(typeof, args)\n    prim = CustomJVPTraced(traced, self.jvp_fun, in_avals, self.symz,\n                           self.static_argnums)\n    return prim(*args)\n\n\nclass MappingSpec: pass\nclass HiPspec:\n  def to_lo(self) -> tuple[PartitionSpec, ...]:\n    _must_override(self, \"to_lo\", \"shard_map\")\n  def to_tangent_spec(self) -> HiPspec:\n    _must_override(self, \"to_tangent_spec\", \"autodiff through shard_map\")\n  def to_ct_spec(self) -> HiPspec:\n    _must_override(self, \"to_ct_spec\", \"autodiff through shard_map\")","sourceCodeStart":1204,"sourceCodeEnd":1240,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L1204-L1240","documentation":"After tracing the custom_jvp-decorated function, JAX inspects the trace's constants (closed-over values). If any constant is a Tracer — meaning the function closes over a value created in an outer jit/grad/vmap scope — it raises UnexpectedTracerError telling you to pass the value as an explicit input instead.","triggerScenarios":"Defining the decorated function inside another transformed function so it captures a traced array via closure, e.g. def make(x_traced): @jax.custom_jvp def g(y): return y * x_traced ... then calling g inside jit.","commonSituations":"Factory/partial-application patterns inside jit or vmap; capturing batched values from vmap in closures; moving code into loops where a closure accidentally captures loop-carried tracers.","solutions":["Rewrite the function to take the closed-over tracer as an explicit argument","Construct the decorated function outside any transformation, passing captured values as inputs at call time","Use functools.partial with positional args instead of closures over traced arrays"],"exampleFix":"# before\ndef make(w):\n  @jax.custom_jvp\n  def g(x):\n    return x * w          # closes over traced w\n  return g\njit(lambda x: make(w_traced)(x))(x)\n# after\n@jax.custom_jvp\ndef g(x, w):\n  return x * w\n@g.defjvp\ndef g_jvp(p, t):\n  (x, w), (xd, wd) = p, t\n  return x * w, xd * w + x * wd\njit(lambda x: g(x, w_traced))(x)","handlingStrategy":"type-guard","validationCode":"# after tracing, ensure no tracer constants:\ntraced = jax.make_jaxpr(g)(y)  # if g closes over outer tracers this leaks/errs\n# guard: define decorated functions at module scope, pass captured arrays as args","typeGuard":"from jax.core import Tracer\ndef closes_over_tracer(fun) -> bool:\n    return any(isinstance(c, Tracer) for c in fun.__closure__ or () if hasattr(c, 'cell_contents'))","tryCatchPattern":null,"preventionTips":["Never define custom_jvp functions inside jit/vmap/grad bodies","Pass captured values as explicit arguments","Use functools.partial on module-scope decorated functions"],"tags":["jax","custom-jvp","tracer-error","closure"],"backgroundTag":"jax-unexpected-tracer-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}