jax-ml/jax · error · UnexpectedTracerError

custom_jvp-decorated function {self.f} closed over a {type(t

Error message

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.

What it means

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.

Source

Thrown at jax/_src/hijax.py:1222

          "The input arguments to the custom_jvp-decorated function "
          f"{self.f.__name__} could not be resolved to positional-only "
          f"arguments. Binding failed with the error:\n{e}") from e
    if any(isinstance(args[i], core.Tracer) for i in self.static_argnums):
      raise UnexpectedTracerError("custom_jvp inputs marked with nondiff_argnums "
                                  "must be static, not Tracers")
    if all(is_hashable(args[i]) for i in self.static_argnums):
      traced = api.jit(self.f, static_argnums=(*self.static_argnums,)).trace(*args)
    else:
      # jit requires hashable static_argnums values, but classic custom_jvp
      # accepted unhashable nondiff_argnums values, so close over them instead
      which_static = [i in self.static_argnums for i in range(len(args))]
      dyn_args, static_args = partition_list(which_static, args)
      f = dyn_args_fun(self.f, self.static_argnums,
                       tuple(map(WrapHashably, static_args)), len(args))
      traced = api.jit(f).trace(*dyn_args)
    if any(isinstance(x, core.Tracer) for x in traced._consts):
      t = next(x for x in traced._consts if isinstance(x, core.Tracer))
      raise UnexpectedTracerError(
          f"custom_jvp-decorated function {self.f} closed over a {type(t).__name__} "
          f"of type {t.aval.str_short()}, but custom_jvp functions can't close "
          f"over Tracers. Rewrite {self.f} to take it as an explicit input.")
    args = tuple(Static(x) if i in self.static_argnums else x for i, x in enumerate(args))
    in_avals = tree_map(typeof, args)
    prim = CustomJVPTraced(traced, self.jvp_fun, in_avals, self.symz,
                           self.static_argnums)
    return prim(*args)


class MappingSpec: pass
class HiPspec:
  def to_lo(self) -> tuple[PartitionSpec, ...]:
    _must_override(self, "to_lo", "shard_map")
  def to_tangent_spec(self) -> HiPspec:
    _must_override(self, "to_tangent_spec", "autodiff through shard_map")
  def to_ct_spec(self) -> HiPspec:
    _must_override(self, "to_ct_spec", "autodiff through shard_map")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Rewrite the function to take the closed-over tracer as an explicit argument
2. Construct the decorated function outside any transformation, passing captured values as inputs at call time
3. Use functools.partial with positional args instead of closures over traced arrays

Example fix

# before
def make(w):
  @jax.custom_jvp
  def g(x):
    return x * w          # closes over traced w
  return g
jit(lambda x: make(w_traced)(x))(x)
# after
@jax.custom_jvp
def g(x, w):
  return x * w
@g.defjvp
def g_jvp(p, t):
  (x, w), (xd, wd) = p, t
  return x * w, xd * w + x * wd
jit(lambda x: g(x, w_traced))(x)

# after tracing, ensure no tracer constants:
traced = jax.make_jaxpr(g)(y)  # if g closes over outer tracers this leaks/errs
# guard: define decorated functions at module scope, pass captured arrays as args

from jax.core import Tracer
def closes_over_tracer(fun) -> bool:
    return any(isinstance(c, Tracer) for c in fun.__closure__ or () if hasattr(c, 'cell_contents'))

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• custom_jvp inputs marked with nondiff_argnums must be static
• Pure callbacks do not support JVP. Please use `jax.custom_jv
• Can't use ``defjvps`` with ``nondiff_argnums``.
• No JVP defined for custom_jvp function {primal_name} using d
• The input arguments to the custom_jvp-decorated function {pr

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/da1b1c53c851b444. Report an issue: GitHub.