{"record":{"id":"b2ec94ebe5a95caf","repo":"jax-ml/jax","slug":"structure-of-output-returned-by-custom-vmap-rule","errorCode":null,"errorMessage":"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}","messagePattern":"structure of output returned by custom vmap rule \\((.+?)\\) does not match that of original custom-vmapped function\\.\\\\nOriginal output: (.+?)\\\\nRule output: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/custom_batching.py","lineNumber":216,"sourceCode":"\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\n# Like batching.batch except (a) not curried and (b) returns inferred output\n# axes instead of accepting and matching a given spec of output axes. Assumes\n# `f` is pytree-flattened\ndef vmap_unrestricted(f: lu.WrappedFun, *args, in_axes, axis_name, axis_size):\n  axis_data = batching.AxisData(axis_name, axis_size, None, None)\n  tag = core.TraceTag()","sourceCodeStart":198,"sourceCodeEnd":234,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/custom_batching.py#L198-L234","documentation":"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).","triggerScenarios":"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.","commonSituations":"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.","solutions":["Return exactly the same container structure as the original function (wrap/unwrap single outputs the same way)","Re-run the primal function to observe its output structure and mirror it in the rule"],"exampleFix":"# before\nreturn (y,), True\n# after\nreturn y, True","handlingStrategy":"validation","validationCode":"import jax\norig_tree = jax.tree_util.tree_structure(f(*sample_args))\n# in rule tests: assert jax.tree_util.tree_structure(rule_out) == orig_tree","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pin the expected output structure in a test fixture","Single-output rules should return bare arrays, not 1-tuples"],"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"}