jax-ml/jax · error · UnexpectedTracerError

custom_vjp inputs marked with nondiff_argnums must be static

Error message

custom_vjp inputs marked with nondiff_argnums must be static, not Tracers

What it means

Arguments of a @jax.custom_vjp function marked as nondiff_argnums must be static (concrete Python values), but one of them was a core.Tracer — i.e. an abstract/traced value arising from jit/grad/vmap/pmap transformation. JAX raises UnexpectedTracerError because static args are passed as Python constants to the rule functions and cannot contain tracers.

Source

Thrown at jax/_src/hijax.py:949

    self.bwd = bwd
    self.symz = symbolic_zeros
    self.opt_remat = optimize_remat

  def defvjp_with_logs(self, fwd, bwd, *, symbolic_zeros=False,
                       optimize_remat=False):
    self.defvjp(fwd, bwd, symbolic_zeros=symbolic_zeros,
                optimize_remat=optimize_remat)
    self.with_logs = True

  def __call__(self, *args, **kwargs):
    if not self.fwd or not self.bwd:
      msg = f"No VJP defined for custom_vjp function {self.f.__name__} using defvjp."
      raise AttributeError(msg)

    args = resolve_kwargs(self.f, args, kwargs)
    if any(isinstance(l, core.Tracer) for i in self.static_argnums
           for l in tree_leaves(args[i])):
      raise UnexpectedTracerError("custom_vjp 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_vjp
      # 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)
    args = tuple(Static(x) if i in self.static_argnums else x for i, x in enumerate(args))
    consts, traced = traced.with_consts_as_arg()
    fwd_ = update_wrapper(lambda _, __, *args: self.fwd(*args), self.fwd)
    static_argnums = frozenset(i + 2 for i in self.static_argnums)
    in_avals = tree_map(typeof, (consts, (), *args))
    prim = CustomVJPTraced(traced, fwd_, self.bwd, in_avals, self.symz,
                           static_argnums, self.opt_remat, self.with_logs)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Move the traced value out of the nondiff argnum and into a differentiable argument
2. Compute the static value outside the transformation and pass it in as a plain Python int/str
3. If the value must remain dynamic, handle it inside the fwd/bwd rules as an array input instead of a static one

Example fix

// before
@jax.custom_vjp, nondiff_argnums=(1,)
def f(x, n): ...
jit(lambda x, n: f(x, n))(x, jnp.asarray(3))  # traced static arg
// after
jit(lambda x, n: f(x, int(n)))(x, jnp.asarray(3))
# or make n a traced diff-arg and use it inside the rules

from jax.core import Tracer
def safe_static_args(fn, args, static_idx):
    for i in static_idx:
        if isinstance(args[i], Tracer):
            raise TypeError(f'arg {i} is traced; convert to static or make it differentiable')

from jax.core import Tracer
def is_static(v) -> bool:
    return not isinstance(v, Tracer)

When it happens

Trigger: Calling a custom_vjp function with nondiff_argnums while a nondiff argument is produced inside a jax.jit, jax.grad, lax.scan, or vmap transformation (e.g. passing a traced array as the 'static' integer argument).

Common situations: Passing a shape or index computed from traced arrays (e.g. x.shape[0] inside jit works, but array-derived values traced); reusing the same function with and without jit; forgetting that nondiff_argnums indexes positional args only.

Related errors

• custom_jvp inputs marked with nondiff_argnums must be static
• Pure callbacks do not support transpose. Please use `jax.cus
• remat optimization for custom_vjp does not support forward f
• remat optimization for custom_vjp does not support higher-or
• under vmap, the {member_name} of {name} produced an output b

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