jax-ml/jax · error · ValueError

Cluster barriers are not arrays

Error message

Cluster barriers are not arrays

What it means

ClusterBarrierSpec.get_array_aval unconditionally raises 'Cluster barriers are not arrays': cluster barriers are synchronization primitives allocated in SMEM, represented only as refs (get_ref_aval). Requesting an array aval is an API misuse, e.g. by generic code that assumes all specs are arrays.

Source

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

@dataclasses.dataclass(frozen=True, kw_only=True)
class ClusterBarrier:
  collective_axes: tuple[str | tuple[str, ...], ...]
  num_barriers: int | Sequence[int] = 1
  num_arrivals: int = 1
  orders_tensor_core: bool = False
  leader_tracked: bool = False

  def __post_init__(self):
    if (n := self.num_arrivals) < 1:
      raise ValueError(f"Num arrivals must be at least 1, but got {n}")

    if isinstance(self.num_barriers, int):
      object.__setattr__(self, "num_barriers", (self.num_barriers,))
    else:
      object.__setattr__(self, "num_barriers", tuple(self.num_barriers))

  def get_array_aval(self) -> jax_core.ShapedArray:
    raise ValueError("Cluster barriers are not arrays")

  def get_ref_aval(self) -> state.AbstractRef:
    ty = ClusterBarrierType(
        collective_axes=self.collective_axes,
        num_arrivals=self.num_arrivals,
        orders_tensor_core=self.orders_tensor_core,
        leader_tracked=self.leader_tracked,
    )
    return state.AbstractRef(jax_core.ShapedArray(self.num_barriers, ty), SMEM)


@dataclasses.dataclass(frozen=True)
class WGMMAAccumulatorRef:
  shape: tuple[int, int]
  dtype: jnp.dtype = jnp.float32
  _init: Any = state_types.uninitialized

  def get_ref_aval(self) -> state.AbstractRef:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Type-dispatch: use get_ref_aval() for ClusterBarrierSpec
2. Exclude barrier specs from array-aval code paths
3. Pass explicit array specs (ArraySpec-style) where arrays are required

Example fix

// before
avals = [s.get_array_aval() for s in specs]  # raises on cluster barriers
// after
avals = [s.get_ref_aval() if isinstance(s, plgpu.ClusterBarrierSpec) else s.get_array_aval() for s in specs]

if isinstance(spec, plgpu.ClusterBarrierSpec):
    aval = spec.get_ref_aval()
else:
    aval = spec.get_array_aval()

def is_cluster_barrier(spec) -> bool:\n    return isinstance(spec, plgpu.ClusterBarrierSpec)

When it happens

Trigger: Calling get_array_aval() on a ClusterBarrierSpec — commonly from generic code that uniformly extracts array avals from all specs, or by passing a cluster barrier spec as an out_shape.

Common situations: Warp-specialized cluster kernels where barrier specs flow through the same plumbing as array specs; refactors that introduced generic aval extraction.

Related errors

• Barriers are not arrays
• Unsupported block dimension type: {type(bd)} for block shape
• Unsupported pipeline mode: {pipeline_mode}.
• 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/bd52144c899f9097. Report an issue: GitHub.