jax-ml/jax · error · ValueError

Unsupported barrier type: {type(barrier)}

Error message

Unsupported barrier type: {type(barrier)}

What it means

The interpreter's barrier-arrive callback only supports the standard memory.Barrier implementation. Arriving at any other synchronization primitive type (e.g. a named barrier with different semantics) is rejected.

Source

Thrown at jax/_src/pallas/mosaic_gpu/interpret/gpu_callbacks.py:942

  barrier, clock = shared_memory.get_barrier_and_increment_clock(
      barrier_key, thread
  )
  if isinstance(barrier, memory.ClusterBarrier):
    barrier.arrive(
        mesh_location=mesh_location,
        thread=thread,
        clock=clock,
        logging_info=memory.GPULoggingInfo(mesh_location, thread, source_info),
    )
  elif isinstance(barrier, memory.Barrier):
    barrier.arrive(
        thread,
        clock,
        memory.GPULoggingInfo(mesh_location, thread, source_info),
    )
  else:
    raise ValueError(f"Unsupported barrier type: {type(barrier)}")
  return token


# Note that this callback is also used for arriving at cluster barriers.
def call_barrier_arrive(
    token: jax.Array,
    mesh_location: memory.MeshLocation,
    thread: memory.Thread,
    allocation_key_as_array: jax.Array,
    source_info: source_info_util.SourceInfo | None = None,
):
  return callback.io_callback(
      functools.partial(_barrier_arrive, source_info=source_info),
      TOKEN_SHAPE_DTYPE,
      token,
      mesh_location,
      thread,
      allocation_key_as_array,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Restructure the kernel to use plain Barrier objects for arrive semantics
2. Update jax / mosaic plugin so the barrier type is supported in interpret mode
3. Run on device instead of interpret mode while using the exotic barrier
4. Report upstream if the barrier type should be supported

def is_supported_barrier(b) -> bool:
    from jax._src.pallas.mosaic_gpu.interpret import memory
    return isinstance(b, memory.Barrier)

try:
    kernel(x)
except ValueError as e:
    if 'Unsupported barrier type' in str(e):
        swap to a plain Barrier

When it happens

Trigger: Calling barrier_arrive / arriving at a barrier object in interpret mode whose runtime type isn't memory.Barrier (unsupported barrier variant such as cgbarrier or future barrier kinds).

Common situations: Using experimental barrier APIs newly added to Mosaic GPU before interpret-mode support catches up; version mismatch between kernel code and jax interpreter.

Related errors

• tcgen05_mma only allows arriving on a Barrier
• tcgen05_commit_arrive only allows arriving on a Barrier
• `thread_id` must be zero when allocating a buffer for all th
• `block_id` must be zero when allocating a buffer for all thr
• Out-of-bounds read of {allocation_key}: reading [{read_range

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