jax-ml/jax · error · UnexpectedTracerError

custom_vjp inputs marked with nondiff_argnums must be…

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.

Solutions

Move the traced value out of the nondiff argnum and into a differentiable argument
Compute the static value outside the transformation and pass it in as a plain Python int/str
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

Defensive patterns

Strategy: type-guard

Validate before calling

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')

Type guard

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

Prevention

Convert traced scalars with int()/float() outside jit before passing as nondiff args
Keep nondiff args to plain Python scalars/strings
Prefer passing arrays as diff args

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

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

Appendix: 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)