jax-ml/jax · error · NotImplementedError

pbroadcast batcher only supports a single axis

Error message

pbroadcast batcher only supports a single axis

What it means

When vmap batches a pbroadcast call, the batcher can only handle a single named axis. If the pbroadcast was declared over multiple axis names and the batched axis is one of them, remaining axes make the transform ambiguous and it raises NotImplementedError.

Source

Thrown at jax/_src/lax/parallel.py:1364

batching.fancy_primitive_batchers[precv_p] = _ppermute_batcher

def _pbroadcast_transpose_rule(t, x, source, axis_name):
  is_source = axis_index(axis_name) == source
  tsum = psum(t, axis_name)
  return [lax.select(is_source, lax.full_like(t, tsum), lax.full_like(t, 0))]

def _pbroadcast_batcher(axis_data, vals_in, dims_in, axis_name, source):
  axis_size = axis_data.size
  (v,), (d,) = vals_in, dims_in
  if not isinstance(axis_name, (tuple, list)):
    axis_name = (axis_name,)
  if d is None and axis_data.name not in axis_name:
    return pbroadcast_p.bind(v, axis_name=axis_name, source=source), None
  if axis_data.name not in axis_name:
    return pbroadcast_p.bind(v, axis_name=axis_name, source=source), d
  remaining_axes = tuple(axis for axis in axis_name if axis != axis_data.name)
  if remaining_axes:
    raise NotImplementedError("pbroadcast batcher only supports a single axis")
  assert axis_name[0] == axis_data.name, "pbroadcast batcher called with a wrong axis!"
  assert source >= 0 and source < axis_size, "collective broadcast doesn't fit in the axis size!"
  if axis_size == 1 and remaining_axes:
    return pbroadcast_p.bind(v, source=source, axis_name=remaining_axes), d
  if d is None:
    return v, d
  return v.take([source] * axis_size, d), d

def _pbroadcast_lowering(ctx, x, *, axis_name, source):
  replica_groups = _replica_groups(ctx.module_context.axis_context, axis_name, None)
  def source_to_front(group):
    return [group[source]] + list(group[:source]) + list(group[source + 1:])
  replica_groups = [source_to_front(group) for group in replica_groups]
  is_spmd = isinstance(
      ctx.module_context.axis_context,
      (SPMDAxisContext, ShardingContext),
  )
  if is_spmd:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Split pbroadcast into one call per axis name
Restructure so pbroadcast is only under shard_map, not vmap
Reduce to a single axis_name before batching

Example fix

// before
lax.pbroadcast(x, axis_name=('a','b'), source=0)  # then vmap over 'a'
// after
x = lax.pbroadcast(x, axis_name='b', source=0)
x = lax.pbroadcast(x, axis_name='a', source=0)

Defensive patterns

Strategy: validation

Validate before calling

def check_pbroadcast_axes(axis_name):
    if isinstance(axis_name, (tuple, list)) and len(axis_name) > 1:
        raise ValueError('use a single axis_name when batching pbroadcast with vmap')

Prevention

Prefer single-axis pbroadcast calls composed together
Keep pbroadcast under shard_map when multiple axes are involved

When it happens

Trigger: jax.vmap(..., axis_name='a') over a lax.pbroadcast(x, axis_name=('a','b'), source=...) call where both axes are named.

Common situations: Nesting pbroadcast over multi-axis meshes and then applying another vmap; combining two vmapped dimensions with a two-name pbroadcast.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

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