jax-ml/jax · error · ValueError

`collective_axes` must be specified when `leader_tracked` is

Error message

`collective_axes` must be specified when `leader_tracked` is set

What it means

copy_gmem_to_smem's leader_tracked parameter marks a copy partitioned across a collective (e.g. a TP or multi-GPU collective axis). To track which leader performs the copy, the implementation must know the collective axes, so leader_tracked without collective_axes is a user API error.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:1303

  flat_dst_transforms, dst_transforms_treedef = tree_util.tree_flatten(
      dst_transforms
  )
  has_barrier = barrier is not None
  if has_barrier:
    barrier, barrier_transforms = state_primitives.get_ref_and_transforms(
        barrier, None, "copy_gmem_to_smem"
    )
    barrier_operands = [barrier]
  else:
    barrier_transforms = []
    barrier_operands = []
  flat_barrier_transforms, barrier_transforms_treedef = tree_util.tree_flatten(
      barrier_transforms
  )
  if isinstance(collective_axes, str):
    collective_axes = (collective_axes,)
  if leader_tracked is not None and collective_axes is None:
    raise ValueError(
        "`collective_axes` must be specified when `leader_tracked` is set"
    )
  copy_gmem_to_smem_p.bind(
      src,
      dst,
      *barrier_operands,
      *flat_src_transforms,
      *flat_dst_transforms,
      *flat_barrier_transforms,
      *[] if predicate is None else [predicate],
      src_transforms_treedef=src_transforms_treedef,
      dst_transforms_treedef=dst_transforms_treedef,
      barrier_transforms_treedef=barrier_transforms_treedef,
      collective_axes=collective_axes,
      leader_tracked=leader_tracked,
      oob_mode=oob_mode,
      has_barrier=has_barrier,
      has_user_predicate=predicate is not None,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass collective_axes (tuple of axis names or a single string) whenever leader_tracked is set.
If you did not intend collective behavior, drop leader_tracked.

Example fix

# before
copy_gmem_to_smem(src, dst, leader_tracked=CopyPartition.PARTITIONED(0))
# after
copy_gmem_to_smem(src, dst,
    leader_tracked=CopyPartition.PARTITIONED(0), collective_axes=('tp',))

Defensive patterns

Strategy: validation

Validate before calling

if leader_tracked is not None:
    assert collective_axes is not None, 'collective_axes required with leader_tracked'

Prevention

Pass collective_axes and leader_tracked together by convention; wrap them in a small helper function.

When it happens

Trigger: Calling copy_gmem_to_smem(..., leader_tracked=CopyPartition.PARTITIONED(axis)) (or similar) without also passing collective_axes.

Common situations: Writing multi-GPU/collective Pallas kernels and copy-pasting a leader_tracked example while forgetting the collective_axes argument added in the same API revision.

Understand the failure class

Background: "missing required argument" and "the following required arguments were not provided": what required-argument errors mean and how to fix them — this error's family across 20 libraries.

Related errors

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