jax-ml/jax · error · AttributeError

No batching rule defined for custom_vmap function {debug_fun

Error message

No batching rule defined for custom_vmap function {debug_fun.func_name} using def_vmap.

What it means

A function decorated with @jax.custom_vmap must register a batching rule via <fun>.def_vmap(...) before it can be vmapped. Calling the wrapper without a registered vmap_rule raises this AttributeError.

Source

Thrown at jax/_src/custom_batching.py:156

    """
    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))
    assert self.vmap_rule is not None
    debug_rule = api_util.debug_info("custom_vmap rule", self.vmap_rule,
                                     (0, args, args), {})
    out_flat = custom_vmap_p.bind(*consts, *args_flat,
                                  call=closed_call,
                                  rule=ClosedRule(self.vmap_rule,
                                                  debug_rule),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Register a batching rule: call f.def_vmap with a rule mapping (axis_size, in_batched, *args) to (outs, out_batched)
2. Verify def_vmap executes at import time, not inside a conditional
3. Check the custom_vmap docstring example for the exact rule signature

Example fix

# before
@jax.custom_vmap
def f(x): return x * 2
# after
@jax.custom_vmap
def f(x): return x * 2
@f.def_vmap
def rule(axis_size, in_batched, x):
    return x * 2, in_batched

assert getattr(f, 'vmap_rule', None) is not None, 'register f.def_vmap before vmapping'

When it happens

Trigger: Decorating a function with @jax.custom_vmap but never calling fun.def_vmap(rule) and then invoking the function (directly or under jax.vmap).

Common situations: Forgetting the def_vmap step after adding the decorator; reordering code so def_vmap runs lazily/never (e.g. inside an uncalled branch); tutorial code that only shows the decorator half.

Related errors

• structure of output value and output batching specification
• structure of output returned by custom vmap rule ({rule_name
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• {name} wrapped function must be passed at least one argument

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