{"record":{"id":"bd52144c899f9097","repo":"jax-ml/jax","slug":"cluster-barriers-are-not-arrays","errorCode":null,"errorMessage":"Cluster barriers are not arrays","messagePattern":"Cluster barriers are not arrays","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1536,"sourceCode":"@dataclasses.dataclass(frozen=True, kw_only=True)\nclass ClusterBarrier:\n  collective_axes: tuple[str | tuple[str, ...], ...]\n  num_barriers: int | Sequence[int] = 1\n  num_arrivals: int = 1\n  orders_tensor_core: bool = False\n  leader_tracked: bool = False\n\n  def __post_init__(self):\n    if (n := self.num_arrivals) < 1:\n      raise ValueError(f\"Num arrivals must be at least 1, but got {n}\")\n\n    if isinstance(self.num_barriers, int):\n      object.__setattr__(self, \"num_barriers\", (self.num_barriers,))\n    else:\n      object.__setattr__(self, \"num_barriers\", tuple(self.num_barriers))\n\n  def get_array_aval(self) -> jax_core.ShapedArray:\n    raise ValueError(\"Cluster barriers are not arrays\")\n\n  def get_ref_aval(self) -> state.AbstractRef:\n    ty = ClusterBarrierType(\n        collective_axes=self.collective_axes,\n        num_arrivals=self.num_arrivals,\n        orders_tensor_core=self.orders_tensor_core,\n        leader_tracked=self.leader_tracked,\n    )\n    return state.AbstractRef(jax_core.ShapedArray(self.num_barriers, ty), SMEM)\n\n\n@dataclasses.dataclass(frozen=True)\nclass WGMMAAccumulatorRef:\n  shape: tuple[int, int]\n  dtype: jnp.dtype = jnp.float32\n  _init: Any = state_types.uninitialized\n\n  def get_ref_aval(self) -> state.AbstractRef:","sourceCodeStart":1518,"sourceCodeEnd":1554,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1518-L1554","documentation":"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.","triggerScenarios":"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.","commonSituations":"Warp-specialized cluster kernels where barrier specs flow through the same plumbing as array specs; refactors that introduced generic aval extraction.","solutions":["Type-dispatch: use get_ref_aval() for ClusterBarrierSpec","Exclude barrier specs from array-aval code paths","Pass explicit array specs (ArraySpec-style) where arrays are required"],"exampleFix":"// before\navals = [s.get_array_aval() for s in specs]  # raises on cluster barriers\n// after\navals = [s.get_ref_aval() if isinstance(s, plgpu.ClusterBarrierSpec) else s.get_array_aval() for s in specs]","handlingStrategy":"type-guard","validationCode":"if isinstance(spec, plgpu.ClusterBarrierSpec):\n    aval = spec.get_ref_aval()\nelse:\n    aval = spec.get_array_aval()","typeGuard":"def is_cluster_barrier(spec) -> bool:\\n    return isinstance(spec, plgpu.ClusterBarrierSpec)","tryCatchPattern":null,"preventionTips":["Special-case barrier specs in generic aval-extraction code"],"tags":["jax","pallas","barrier","cluster","api-misuse"],"backgroundTag":"wrong-aval-kind-requested","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}