jax-ml/jax · error · TypeError

The input arguments to the custom_vmap-decorated function {d

Error message

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

What it means

jax.custom_vmap-decorated functions must be callable with positional arguments only. This error wraps the TypeError raised when Python cannot bind the supplied keyword arguments to the function signature (unknown kwarg, or kwargs that the flattener cannot reorder).

Source

Thrown at jax/_src/custom_batching.py:149

        and (3) the batched arguments. It should return a tuple of the batched
        output and a pytree of booleans with the same structure as the output,
        specifying whether each output element is batched. See the documentation
        for :py:func:`jax.custom_batching.custom_vmap` for some examples.

    Returns:
      This method passes the rule through, returning ``vmap_rule`` unchanged.
    """
    self.vmap_rule = vmap_rule
    return vmap_rule

  @traceback_util.api_boundary
  def __call__(self, *args, **kwargs):
    debug_fun = api_util.debug_info("custom_vmap fun", self.fun,
                                    args, kwargs)
    try:
      args = api_util.resolve_kwargs(self.fun, args, kwargs)
    except TypeError as e:
      raise TypeError(
          "The input arguments to the custom_vmap-decorated function "
          f"{debug_fun.func_name} could not be resolved to positional-only "
          f"arguments. Binding failed with the error:\n{e}"
      ) from e

    if not self.vmap_rule:
      raise AttributeError(
          f"No batching rule defined for custom_vmap function {debug_fun.func_name} "
          "using def_vmap.")
    args_flat, in_tree = tracing_registry.flatten(args)
    in_avals = [core.typeof(x) for x in args_flat]
    jaxpr, out_avals = pe.trace_to_jaxpr(
        self.fun, ft.pack((ft.treedef_args_to_ft(in_tree, in_avals), {})),
        debug_fun)
    closed_call, consts = pe.separate_consts(jaxpr)
    out_tree = tree_structure(out_avals.unflatten())
    in_tree = treedef_tuple_tracing_registry(
        (tracing_registry.flatten(consts)[1], in_tree))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Call the function with all arguments positionally
2. If keyword usage is desired, wrap the function to bind kwargs first or accept **kwargs explicitly and delegate positionally

Example fix

# before
out = f(x, y, scale=2.0)
# after
out = f(x, y, 2.0)

import inspect
def call_positional(fn, *args, **kwargs):
    sig = inspect.signature(fn)
    return fn(*sig.bind(None, *args, **kwargs).args) if kwargs else fn(*args)

When it happens

Trigger: Calling a @jax.custom_vmap function with keyword arguments, e.g. f(x, scale=2.0) where scale is a regular parameter, or passing a kwarg name not present in the signature.

Common situations: Refactoring a function signature (renaming params) while callers still use old kwarg names; applying vmap/jit pipelines that forward **kwargs into custom_vmap functions.

Related errors

• The input arguments to the custom_jvp-decorated function {pr
• No batching rule defined for custom_vmap function {debug_fun
• structure of output value and output batching specification
• structure of output returned by custom vmap rule ({rule_name
• keyword arguments could not be resolved to positions

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