jax-ml/jax · error · ValueError

structure of output returned by custom vmap rule ({rule_name

Error message

structure of output returned by custom vmap rule ({rule_name(rule)}) does not match that of original custom-vmapped function.\nOriginal output: {original_out_tree}\nRule output: {out_tree}

What it means

The output pytree structure returned by a custom vmap rule must match the structure the original custom-vmapped function produces. This catches rules that change the shape of the output container (e.g. returning a list instead of a scalar).

Source

Thrown at jax/_src/custom_batching.py:216

def ensure_list(xs):
  return xs if type(xs) is list else list(xs)

def rule_name(rule):
  return getattr(rule, '__name__', '<unnamed rule>')

def call_rule(rule, axis_size, in_batched, args):
  return rule(axis_size, ensure_list(in_batched), *args)

def check_vmap_rule_trees(rule, original_out_tree, out_tree, out_batched_tree):
  if out_tree != out_batched_tree:
    raise ValueError(
        'structure of output value and output batching specification returned '
        f'by custom vmap rule ({rule_name(rule)}) do not match.\n'
        f'Output values: {out_tree}\n'
        f'Batching spec: {out_batched_tree}')
  if out_tree != original_out_tree:
    raise ValueError(
        f'structure of output returned by custom vmap rule ({rule_name(rule)}) '
        'does not match that of original custom-vmapped function.\n'
        f'Original output: {original_out_tree}\n'
        f'Rule output: {out_tree}')

# Like batching.bdim_at_front, but doesn't broadcast if not mapped
def maybe_bdim_at_front(x, bdim):
  if bdim is None:
    return x
  else:
    return util.moveaxis(x, bdim, 0)

# Like batching.batch except (a) not curried and (b) returns inferred output
# axes instead of accepting and matching a given spec of output axes. Assumes
# `f` is pytree-flattened
def vmap_unrestricted(f: lu.WrappedFun, *args, in_axes, axis_name, axis_size):
  axis_data = batching.AxisData(axis_name, axis_size, None, None)
  tag = core.TraceTag()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Return exactly the same container structure as the original function (wrap/unwrap single outputs the same way)
2. Re-run the primal function to observe its output structure and mirror it in the rule

Example fix

# before
return (y,), True
# after
return y, True

import jax
orig_tree = jax.tree_util.tree_structure(f(*sample_args))
# in rule tests: assert jax.tree_util.tree_structure(rule_out) == orig_tree

When it happens

Trigger: A def_vmap rule that returns outputs wrapped/unwrapped differently than the primal function — e.g. primal returns a bare array, rule returns a 1-element tuple.

Common situations: Copying a rule from a multi-output function to a single-output one; incremental edits to output packing in one place but not the other.

Related errors

• structure of output value and output batching specification
• No batching rule defined for custom_vmap function {debug_fun
• numpy masked arrays are not supported as direct inputs to JA
• Python int {value} too large to convert to int64
• Python int {value} too large to convert to int32

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