jax-ml/jax · error · TypeError

The input arguments to the custom_jvp-decorated function

Error message

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

What it means

custom_jjp call machinery must rebind the original call's arguments to positional form before tracing, because the flattening transformation is positional-only. inspect-style binding of the caller's args/kwargs against the wrapped function's signature failed, typically because keywords don't match the signature or required static (nondiff) args were supplied by keyword.

Solutions

1. Call with all arguments positionally
2. If keyword ergonomics matter, add an outer plain-Python wrapper that binds kwargs and calls the custom_jvp function positionally

Example fix

# before
loss = f(x, y, eps=1e-3)
# after
loss = f(x, y, 1e-3)

import inspect
def call_positional(fn, *args, **kwargs):
    b = inspect.signature(fn).bind(*args, **kwargs)
    b.apply_defaults()
    return fn(*b.args)

When it happens

Trigger: Calling a @jax.custom_jvp function with keyword arguments that cannot be resolved to positions (unknown names, keyword-only params, or **kwargs forwarding).

Common situations: Refactored parameter names with stale callers; passing config flags as kwargs into a custom_jvp-wrapped loss; wrappers that forward **user_kwargs blindly.

Related errors

• The input arguments to the custom_jvp-decorated function
• The input arguments to the custom_vmap-decorated function
• Can't use ``defjvps`` with ``nondiff_argnums``.
• Can't use ``defjvps`` with ``nondiff_argnums``.
• custom_jvp-decorated function

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

Appendix: source

Thrown at jax/_src/custom_derivatives.py:268

      tangent_out = tree_map(_sum_tangents, primal_out, *all_tangents_out)
      return primal_out, tangent_out

    self.defjvp(jvp)

  @partial(traceback_util.api_boundary,
           repro_api_name="jax.custom_jvp.__call__")
  def __call__(self, *args: Any, **kwargs: Any) -> ReturnValue:
    debug = debug_info("custom_jvp fun", self.fun, args, kwargs,
                       static_argnums=self.nondiff_argnums)
    primal_name = debug.func_name
    if not self.jvp:
      msg = f"No JVP defined for custom_jvp function {primal_name} using defjvp."
      raise AttributeError(msg)

    try:
      args = resolve_kwargs(self.fun, args, kwargs)
    except TypeError as e:
      raise TypeError(
          "The input arguments to the custom_jvp-decorated function "
          f"{primal_name} could not be resolved to positional-only arguments. "
          f"Binding failed with the error:\n{e}"
      ) from e

    if self.nondiff_argnums:
      args = tuple(_stop_gradient(x) if i in self.nondiff_argnums else x
                   for i, x in enumerate(args))
      diff_argnums = [i for i in range(len(args)) if i not in self.nondiff_argnums]
      f_, dyn_args = argnums_partial(lu.wrap_init(self.fun, debug_info=debug),
                                     diff_argnums, args,
                                     require_static_args_hashable=False)
      static_args = [args[i] for i in self.nondiff_argnums]
      diff_args = [args[i] for i, a in enumerate(args) if i not in self.nondiff_argnums]
      debug_jvp = debug_info("custom_jvp jvp", self.jvp,
                             (*static_args, diff_args, diff_args),
                             {},
                             static_argnums=tuple(range(len(static_args))))

View on GitHub (pinned to 1e1c6a8fc0)