{"record":{"id":"d3af7786e47daabf","repo":"jax-ml/jax","slug":"structure-of-output-value-and-output-batching-spec","errorCode":null,"errorMessage":"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}","messagePattern":"structure of output value and output batching specification returned by custom vmap rule \\((.+?)\\) do not match\\.\\\\nOutput values: (.+?)\\\\nBatching spec: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_batching.py","lineNumber":210,"sourceCode":"    consts_batched, in_batched = all_in_batched\n    assert not any(tree_util.tree_leaves(consts_batched)), consts_batched\n    return call_rule(self.rule, axis_size, in_batched, args)\n\n  def __str__(self):\n    return str(self.rule)\n\ndef ensure_list(xs):\n  return xs if type(xs) is list else list(xs)\n\ndef rule_name(rule):\n  return getattr(rule, '__name__', '<unnamed rule>')\n\ndef call_rule(rule, axis_size, in_batched, args):\n  return rule(axis_size, ensure_list(in_batched), *args)\n\ndef check_vmap_rule_trees(rule, original_out_tree, out_tree, out_batched_tree):\n  if out_tree != out_batched_tree:\n    raise ValueError(\n        'structure of output value and output batching specification returned '\n        f'by custom vmap rule ({rule_name(rule)}) do not match.\\n'\n        f'Output values: {out_tree}\\n'\n        f'Batching spec: {out_batched_tree}')\n  if out_tree != original_out_tree:\n    raise ValueError(\n        f'structure of output returned by custom vmap rule ({rule_name(rule)}) '\n        'does not match that of original custom-vmapped function.\\n'\n        f'Original output: {original_out_tree}\\n'\n        f'Rule output: {out_tree}')\n\n# Like batching.bdim_at_front, but doesn't broadcast if not mapped\ndef maybe_bdim_at_front(x, bdim):\n  if bdim is None:\n    return x\n  else:\n    return util.moveaxis(x, bdim, 0)\n","sourceCodeStart":192,"sourceCodeEnd":228,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_batching.py#L192-L228","documentation":"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.","triggerScenarios":"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.","commonSituations":"Returning (out, out_batched) where out is a tuple but out_batched is a scalar bool; forgetting to wrap singleton outputs/specs consistently.","solutions":["Make the batching spec mirror the output structure exactly (tuple of bools for tuple outputs, single bool for single array)","Unpack/flatten outputs and specs through the same pytree shape before returning","Write a quick unit test that runs jax.vmap over the function at import"],"exampleFix":"# before\nreturn (y1, y2), [True]\n# after\nreturn (y1, y2), (True, True)","handlingStrategy":"type-guard","validationCode":null,"typeGuard":"import jax\n\ndef spec_matches(out, batched) -> bool:\n    try:\n        jax.tree_util.tree_structure(out) == jax.tree_util.tree_structure(batched)\n        return True\n    except Exception:\n        return False","tryCatchPattern":null,"preventionTips":["Derive out_batched with tree_map(lambda _: True, outs) then override","Test rules with jax.vmap directly"],"tags":["jax","custom-vmap","pytree","vmap"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}