jax-ml/jax · error · ValueError

collective_axes is not supported in pallas_call. Use plgpu.k

Error message

collective_axes is not supported in pallas_call. Use plgpu.kernel with plgpu.emit_pipeline_warp_specialized instead.

What it means

Raised by the pallas_call-interpreting interpreter's to_block_mapping: the collective_axes option (used by warp-specialized collective pipelines) is not implemented for pallas_call. The message points users to plgpu.kernel with plgpu.emit_pipeline_warp_specialized instead.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:1423

  transforms: Sequence[state_types.Transform] = ()
  delay_release: int = 0
  collective_axes: tuple[Hashable, ...] | None = None
  oob_fill_mode: OOBFillMode = OOBFillMode.ZEROS

  def to_block_mapping(
      self,
      origin: pallas_core.OriginStr,
      array_aval: jax_core.ShapedArray,
      *,
      index_map_avals: Sequence[jax_core.AbstractValue],
      index_map_tree: tree_util.PyTreeDef,
      grid: pallas_core.GridMappingGrid,
      vmapped_dims: tuple[int, ...],
      allow_captured_consts: bool = False,
      debug: bool = False,
  ) -> pallas_core.BlockMapping:
    if self.collective_axes:
      raise ValueError(
          "collective_axes is not supported in pallas_call. Use plgpu.kernel"
          " with plgpu.emit_pipeline_warp_specialized instead."
      )
    bm = super().to_block_mapping(
        origin,
        array_aval,
        index_map_avals=index_map_avals,
        index_map_tree=index_map_tree,
        grid=grid,
        vmapped_dims=vmapped_dims,
        allow_captured_consts=allow_captured_consts,
        debug=debug,
    )
    block_inner_aval = bm.block_aval.inner_aval
    for t in self.transforms:
      block_inner_aval = t.transform_type(block_inner_aval)
    return bm.replace(
        transformed_block_aval=bm.block_aval.update(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Switch the kernel to plgpu.kernel with plgpu.emit_pipeline_warp_specialized, which supports collective_axes
2. Remove/clear collective_axes if you don't need cluster collectives in pallas_call
3. Restructure the collective as an explicit distributed matmul pattern supported by pallas_call

Example fix

// before
pl.pallas_call(kernel, out_shape, collective_axes=(0,))(...)  # ValueError
// after
plgpu.kernel(kernel, out_shape, interpreter=plgpu.emit_pipeline_warp_specialized)(...)

if getattr(config, 'collective_axes', None):
    raise SystemExit('collective_axes requires plgpu.kernel + emit_pipeline_warp_specialized')

When it happens

Trigger: Passing an interpreter/config with collective_axes set (non-empty) to the Mosaic pallas_call path — e.g. a ClusterBarrierType or kernel spec with collective_axes used from pallas_call.

Common situations: Porting plgpu.kernel pipeline code to pallas_call; setting collective_axes (for cluster collectives like collective matmuls) in a kernel invoked through pallas_call.

Related errors

• Only collective_axes that include all JAX device mesh axes
• Expected exactly one collective axis, got {collective_axes=}
• Only the TMA implementation supports collective copies
• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f

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