jax-ml/jax · error · ValueError

under vmap, the {member_name} of {name} produced an output b

Error message

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

What it means

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).

Source

Thrown at jax/_src/hijax.py:335

  def batch_dim_rule(self, axis_data, in_dims):
    fix = lambda d, d_: d if (d is None or d_ is None) else d - (d_ < d)
    in_dims_ = tree_map(fix, in_dims, self.in_dims, is_leaf=lambda x: x is None)
    out_dim = self.prim.batch_dim_rule(axis_data, in_dims_)  # pyrefly: ignore[missing-attribute]
    unfix = lambda d, d_: d if (d is None or d_ is None) else d + (d_ < d)
    return tree_map(unfix, out_dim, self.out_dim, is_leaf=lambda x: x is None)

@contextmanager
def _explain_overbatched_member(prim, member_name):
  try:
    yield
  except ValueError as e:
    if ('but output was batched' not in str(e) and
        'vmap has mapped output' not in str(e)):
      raise
    name = getattr(getattr(prim, 'traced', None), 'fun_name',
                   type(prim).__name__)
    raise ValueError(
        f"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 "
        "`batch_dim_rule` (or `batch`) method to declare the batched "
        "outputs.") from e

def map_zero(axis_data, d, ct):
  if isinstance(ct, ad_util.Zero):
    return ad_util.Zero(core.mapped_aval(axis_data.size, d, ct.aval))
  return ct

def unmap_zero(axis_data, d, ct):
  if isinstance(ct, ad_util.Zero):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Restructure the rule so outputs match the primal's batchedness (e.g. avoid introducing the mapped axis in tangents)
  2. Define the operation as a jax.experimental.hijax.HiPrim and override the relevant member so batchedness is explicit
  3. Move the vmap inside the custom function so the rule sees already-batched inputs

Example fix

# before
@custom_vjp
def f(x, y): ...
# after: make batchedness explicit
from jax.experimental import hijax
class F(hijax.HiPrim):
  ...
Defensive patterns

Strategy: fallback

Try / catch

try:
    jax.vmap(jax.grad(f))(xs)
except ValueError as e:
    if 'inferred to be unbatched' in str(e):
        return jax.vmap(lambda x: jax.grad(f)(x))(xs)  # vmap inside instead
    raise

Prevention

When it happens

Trigger: 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.

Common situations: Gradient of a masked or gather-style op under vmap; a backward rule that broadcasts or expands along the mapped axis.

Related errors

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