jax-ml/jax · error · ValueError

Barriers are not arrays

Error message

Barriers are not arrays

What it means

BarrierSpec.get_array_aval always raises because a BarrierSpec describes hardware synchronization barriers, not a data array. Barriers have no array semantics; only get_ref_aval (an AbstractRef in SMEM) is meaningful. Calling the array path signals a programming/API misuse.

Source

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

      guarantee that the TensorCore-related operations in other threads have
      completed. Similarly, when False any TensorCore operation in the waiting
      thread is allowed to begin before the wait succeeds.
  """
  num_arrivals: int = 1
  num_barriers: int | Sequence[int] = 1
  orders_tensor_core: 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("Barriers are not arrays")

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


@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}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Branch on the spec type: use get_ref_aval() for BarrierSpec/ClusterBarrierSpec, get_array_aval() only for array specs
2. Don't pass BarrierSpec where an out_shape/array is required
3. Keep barriers out of any code path that computes array shapes/dtypes

Example fix

// before
aval = spec.get_array_aval()  # raises for barriers
// after
aval = spec.get_ref_aval() if isinstance(spec, plgpu.BarrierSpec) else spec.get_array_aval()

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

def is_barrier_spec(spec) -> bool:\n    return isinstance(spec, plgpu.BarrierSpec)

When it happens

Trigger: Code that treats a BarrierSpec as an output/input array — e.g. passing a barrier spec where an out_shape/array aval is expected, or generic code that calls get_array_aval() on every spec.

Common situations: Generic plumbing in user kernels or JAX transforms that iterates specs and requests array avals unconditionally; confusing barrier buffers with regular SMEM arrays.

Related errors

• Cluster 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/c5912c5345778d2b. Report an issue: GitHub.