jax-ml/jax · error · ValueError

structure of output value and output batching specification

Error message

structure of output value and output batching specification returned by custom vmap rule ({rule_name(rule)}) do not match.\nOutput values: {out_tree}\nBatching spec: {out_batched_tree}

What it means

A custom vmap rule must return two parallel structures: the output values and a batching specification (which outputs carry the batch axis). This ValueError fires when those two pytrees differ in structure.

Source

Thrown at jax/_src/custom_batching.py:210

    consts_batched, in_batched = all_in_batched
    assert not any(tree_util.tree_leaves(consts_batched)), consts_batched
    return call_rule(self.rule, axis_size, in_batched, args)

  def __str__(self):
    return str(self.rule)

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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the batching spec mirror the output structure exactly (tuple of bools for tuple outputs, single bool for single array)
2. Unpack/flatten outputs and specs through the same pytree shape before returning
3. Write a quick unit test that runs jax.vmap over the function at import

Example fix

# before
return (y1, y2), [True]
# after
return (y1, y2), (True, True)

import jax

def spec_matches(out, batched) -> bool:
    try:
        jax.tree_util.tree_structure(out) == jax.tree_util.tree_structure(batched)
        return True
    except Exception:
        return False

When it happens

Trigger: Writing a def_vmap rule whose returned out_batched spec (e.g. a list [True]) does not structurally match the returned outputs (e.g. a single array or nested tuple), or vice versa.

Common situations: Returning (out, out_batched) where out is a tuple but out_batched is a scalar bool; forgetting to wrap singleton outputs/specs consistently.

Related errors

• structure of output returned by custom vmap rule ({rule_name
• 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/d3af7786e47daabf. Report an issue: GitHub.