jax-ml/jax · error · UnexpectedTracerError

custom_jvp inputs marked with nondiff_argnums must be static

Error message

custom_jvp inputs marked with nondiff_argnums must be static, not Tracers

What it means

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.

Source

Thrown at jax/_src/hijax.py:1208

      tangent_out = tree_map(sum_tangents, primal_out, *all_tangents_out)
      return primal_out, tangent_out
    self.defjvp(jvp)

  def __call__(self, *args, **kwargs):
    if not self.jvp_fun:
      msg = (f"No JVP defined for custom_jvp function {self.f.__name__} "
             "using defjvp.")
      raise AttributeError(msg)

    try:
      args = resolve_kwargs(self.f, args, kwargs)
    except TypeError as e:
      raise TypeError(
          "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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Convert to a concrete Python value before the call (int(n), .item() outside jit)
2. Move the argument into the differentiable args and handle it inside the rule
3. Hoist computation of the static value outside the transforming function

Example fix

// before
@jax.custom_jvp(nondiff_argnums=(1,))
def f(x, n): ...
jit(lambda x, n: f(x, n))(x, traced_n)
// after
n_static = int(traced_n)  # computed outside jit
jit(lambda x: f(x, n_static))(x)

from jax.core import Tracer
assert not any(isinstance(args[i], Tracer) for i in f.static_argnums)

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

When it happens

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

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

Related errors

• custom_vjp inputs marked with nondiff_argnums must be static
• custom_jvp-decorated function {self.f} closed over a {type(t
• 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

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