jax-ml/jax · error · TypeError

The input arguments to the custom_jvp-decorated function {pr

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

Like custom_vmap, custom_jvp functions must be invoked with positional arguments; this TypeError wraps the failure to bind supplied keyword arguments to positional parameters.

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)

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_vmap-decorated function {d
• The input arguments to the custom_jvp-decorated function {se
• 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/75dc7a40bbf407fd. Report an issue: GitHub.