jax-ml/jax · error · TypeError

The input arguments to the custom_jvp-decorated function {se

Error message

The input arguments to the custom_jvp-decorated function {self.f.__name__} could not be resolved to positional-only arguments. Binding failed with the error:
{e}

What it means

The custom_jvp wrapper resolves keyword arguments to positional-only arguments using the decorated function's signature (resolve_kwargs). If binding fails — wrong kwarg name, unexpected kwargs, or missing required ones — the original TypeError is re-raised with this message explaining the positional-only constraint.

Source

Thrown at jax/_src/hijax.py:1203

      primal_out = self(*primals)
      zeros = tree_map(ad_util.p2tz, primal_out)
      all_tangents_out = [j(t, primal_out, *primals) if j else zeros
                          for t, j in zip(tangents, jvps)]
      sum_tangents = lambda _, x, *xs: reduce(ad.add_tangents, xs, x)
      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))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Call the function with positional arguments only
Fix kwarg names to match the decorated function's signature exactly
When wrapping, normalize kwargs to positional with inspect.signature(...).bind before calling

Example fix

// before
@jax.custom_jvp
def f(x, scale):
  return x * scale
f(x=1.0, scales=2.0)  # typo'd kwarg
// after
f(1.0, 2.0)
# or f(x=1.0, scale=2.0)

Defensive patterns

Strategy: validation

Validate before calling

import inspect
sig = inspect.signature(f)
bound = sig.bind(*args, **kwargs)  # raises early with a clear error if binding fails
f(*bound.args)

Try / catch

try:
    f(**kwargs)
except TypeError as e:
    if 'could not be resolved to positional-only' in str(e):
        # retry with positional args in signature order
        f(*positional_in_order)

Prevention

Call custom_jvp functions with positional arguments
Validate kwargs with inspect.signature.bind before forwarding

When it happens

Trigger: Calling a @custom_jvp function with keyword arguments that don't bind to its signature, e.g. f(x=1, nonexistent=2), or calling with kwargs when the wrapped function uses *args or positional-only params.

Common situations: Refactoring function parameter names while callers still use old kwarg names; wrapper functions that forward **kwargs blindly; methods decorated with custom_jvp where 'self' shifts positions and breaks static_argnums indices.

Related errors

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