{"record":{"id":"16f4af6fe106d966","repo":"jax-ml/jax","slug":"under-vmap-the-member-name-of-name-produced-a","errorCode":null,"errorMessage":"under vmap, the {member_name} of {name} produced an output batched along the mapped axis where the application itself was inferred to be unbatched. The batchedness of a custom_jvp/custom_vjp application under vmap is inferred from its primal function alone, but a rule may produce more-batched outputs (e.g. if a tangent depends on a batched input that the primal output does not use). To support that, define the operation as a jax.experimental.hijax.HiPrim and override its ","messagePattern":"under vmap, the (.+?) of (.+?) produced an output batched along the mapped axis where the application itself was inferred to be unbatched\\. The batchedness of a custom_jvp/custom_vjp application under vmap is inferred from its primal function alone, but a rule may produce more-batched outputs \\(e\\.g\\. if a tangent depends on a batched input that the primal output does not use\\)\\. To support that, define the operation as a jax\\.experimental\\.hijax\\.HiPrim and override its ","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/hijax.py","lineNumber":335,"sourceCode":"\n  def batch_dim_rule(self, axis_data, in_dims):\n    fix = lambda d, d_: d if (d is None or d_ is None) else d - (d_ < d)\n    in_dims_ = tree_map(fix, in_dims, self.in_dims, is_leaf=lambda x: x is None)\n    out_dim = self.prim.batch_dim_rule(axis_data, in_dims_)  # pyrefly: ignore[missing-attribute]\n    unfix = lambda d, d_: d if (d is None or d_ is None) else d + (d_ < d)\n    return tree_map(unfix, out_dim, self.out_dim, is_leaf=lambda x: x is None)\n\n@contextmanager\ndef _explain_overbatched_member(prim, member_name):\n  try:\n    yield\n  except ValueError as e:\n    if ('but output was batched' not in str(e) and\n        'vmap has mapped output' not in str(e)):\n      raise\n    name = getattr(getattr(prim, 'traced', None), 'fun_name',\n                   type(prim).__name__)\n    raise ValueError(\n        f\"under vmap, the {member_name} of {name} produced an output batched \"\n        \"along the mapped axis where the application itself was inferred to \"\n        \"be unbatched. The batchedness of a custom_jvp/custom_vjp \"\n        \"application under vmap is inferred from its primal function alone, \"\n        \"but a rule may produce more-batched outputs (e.g. if a tangent \"\n        \"depends on a batched input that the primal output does not use). \"\n        \"To support that, define the operation as a \"\n        \"jax.experimental.hijax.HiPrim and override its \"\n        \"`batch_dim_rule` (or `batch`) method to declare the batched \"\n        \"outputs.\") from e\n\ndef map_zero(axis_data, d, ct):\n  if isinstance(ct, ad_util.Zero):\n    return ad_util.Zero(core.mapped_aval(axis_data.size, d, ct.aval))\n  return ct\n\ndef unmap_zero(axis_data, d, ct):\n  if isinstance(ct, ad_util.Zero):","sourceCodeStart":317,"sourceCodeEnd":353,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/hijax.py#L317-L353","documentation":"Under vmap, a custom_jvp/custom_vjp rule produced output batched along the mapped axis although the primal application was inferred unbatched. JAX infers batchedness of custom-rule applications from the primal only, so over-batched rule output triggers this wrapper error (suggesting HiPrim with overridden members).","triggerScenarios":"jax.vmap around a custom_jvp/custom_vjp function where a tangent/cotangent rule depends on a batched input the primal output does not use, raising an inner 'output was batched' / 'vmap has mapped output' ValueError which is re-raised with this explanation.","commonSituations":"Gradient of a masked or gather-style op under vmap; a backward rule that broadcasts or expands along the mapped axis.","solutions":["Restructure the rule so outputs match the primal's batchedness (e.g. avoid introducing the mapped axis in tangents)","Define the operation as a jax.experimental.hijax.HiPrim and override the relevant member so batchedness is explicit","Move the vmap inside the custom function so the rule sees already-batched inputs"],"exampleFix":"# before\n@custom_vjp\ndef f(x, y): ...\n# after: make batchedness explicit\nfrom jax.experimental import hijax\nclass F(hijax.HiPrim):\n  ...\n","handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    jax.vmap(jax.grad(f))(xs)\nexcept ValueError as e:\n    if 'inferred to be unbatched' in str(e):\n        return jax.vmap(lambda x: jax.grad(f)(x))(xs)  # vmap inside instead\n    raise","preventionTips":["Keep rule output batchedness consistent with the primal","Consider HiPrim when rules may over-batch","Test grads under vmap early in development"],"tags":["jax","vmap","custom-vjp","custom-jvp","batching"],"backgroundTag":"vmap-batchedness-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}