{"record":{"id":"d078d3da75a8d088","repo":"jax-ml/jax","slug":"pbroadcast-batcher-only-supports-a-single-axis","errorCode":null,"errorMessage":"pbroadcast batcher only supports a single axis","messagePattern":"pbroadcast batcher only supports a single axis","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":1364,"sourceCode":"batching.fancy_primitive_batchers[precv_p] = _ppermute_batcher\n\ndef _pbroadcast_transpose_rule(t, x, source, axis_name):\n  is_source = axis_index(axis_name) == source\n  tsum = psum(t, axis_name)\n  return [lax.select(is_source, lax.full_like(t, tsum), lax.full_like(t, 0))]\n\ndef _pbroadcast_batcher(axis_data, vals_in, dims_in, axis_name, source):\n  axis_size = axis_data.size\n  (v,), (d,) = vals_in, dims_in\n  if not isinstance(axis_name, (tuple, list)):\n    axis_name = (axis_name,)\n  if d is None and axis_data.name not in axis_name:\n    return pbroadcast_p.bind(v, axis_name=axis_name, source=source), None\n  if axis_data.name not in axis_name:\n    return pbroadcast_p.bind(v, axis_name=axis_name, source=source), d\n  remaining_axes = tuple(axis for axis in axis_name if axis != axis_data.name)\n  if remaining_axes:\n    raise NotImplementedError(\"pbroadcast batcher only supports a single axis\")\n  assert axis_name[0] == axis_data.name, \"pbroadcast batcher called with a wrong axis!\"\n  assert source >= 0 and source < axis_size, \"collective broadcast doesn't fit in the axis size!\"\n  if axis_size == 1 and remaining_axes:\n    return pbroadcast_p.bind(v, source=source, axis_name=remaining_axes), d\n  if d is None:\n    return v, d\n  return v.take([source] * axis_size, d), d\n\ndef _pbroadcast_lowering(ctx, x, *, axis_name, source):\n  replica_groups = _replica_groups(ctx.module_context.axis_context, axis_name, None)\n  def source_to_front(group):\n    return [group[source]] + list(group[:source]) + list(group[source + 1:])\n  replica_groups = [source_to_front(group) for group in replica_groups]\n  is_spmd = isinstance(\n      ctx.module_context.axis_context,\n      (SPMDAxisContext, ShardingContext),\n  )\n  if is_spmd:","sourceCodeStart":1346,"sourceCodeEnd":1382,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L1346-L1382","documentation":"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.","triggerScenarios":"jax.vmap(..., axis_name='a') over a lax.pbroadcast(x, axis_name=('a','b'), source=...) call where both axes are named.","commonSituations":"Nesting pbroadcast over multi-axis meshes and then applying another vmap; combining two vmapped dimensions with a two-name pbroadcast.","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"],"exampleFix":"// before\nlax.pbroadcast(x, axis_name=('a','b'), source=0)  # then vmap over 'a'\n// after\nx = lax.pbroadcast(x, axis_name='b', source=0)\nx = lax.pbroadcast(x, axis_name='a', source=0)","handlingStrategy":"validation","validationCode":"def check_pbroadcast_axes(axis_name):\n    if isinstance(axis_name, (tuple, list)) and len(axis_name) > 1:\n        raise ValueError('use a single axis_name when batching pbroadcast with vmap')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Prefer single-axis pbroadcast calls composed together","Keep pbroadcast under shard_map when multiple axes are involved"],"tags":["jax","pbroadcast","vmap","not-implemented"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}