jax-ml/jax · error · ValueError

Barrier does not support arbitrary transforms

Error message

Barrier does not support arbitrary transforms

What it means

_get_barrier_base_index only handles transform lists that begin with an NDIndexer (plain indexing). Any other transform (swizzles, tilings, transposes, compositions) applied to a barrier ref cannot be lowered to an address, so it is rejected.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:1515

            num_int_idxs += 1

          if isinstance(
              idx, (int, ir.Value, mgpu.FragmentedArray, literals.TypedNdArray)
          ):
            idx = lowering._as_index(idx)  # pylint: disable=protected-access
          else:
            raise ValueError(
                "Barrier can only be indexed with integers or slices, got"
                f" {idx}"
            )

          idx = arith_dialect.muli(idx, lowering._as_index(stride))  # pylint: disable=protected-access
          if base_index is None:
            base_index = idx
          else:
            base_index = arith_dialect.addi(base_index, idx)
      case _:
        raise ValueError("Barrier does not support arbitrary transforms")
  return base_index


barrier_arrive_p = jax_core.Primitive("barrier_arrive")
barrier_arrive_p.multiple_results = True


@barrier_arrive_p.def_effectful_abstract_eval
def _barrier_arrive_abstract_eval(barrier, *args, **params):
  del args, params  # Unused.
  _check_ref(barrier, "barrier", gpu_core.SMEM)
  return (), {gpu_core._memory_effect}


def _barrier_arrive_pp_eqn(
    eqn: jax_core.JaxprEqn,
    context: jax_core.JaxprPpContext,
    settings: jax_core.JaxprPpSettings,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove transforms from barrier refs; keep them plain SMEM-allocated buffers.
2. Apply transforms only to data refs, using separate code paths for barriers.

from jax._src.pallas.mosaic_gpu import indexing
assert all(isinstance(t, indexing.NDIndexer) for t in transforms), 'barriers take only indexing transforms'

When it happens

Trigger: Applying non-indexing transforms (e.g. via TransformedRef with swizzle/transpose) to a barrier ref and then arriving/waiting on it.

Common situations: Reusing generic ref-transformation helpers that add layout transforms to all refs in a kernel, accidentally including barriers/accumulators.

Related errors

• Transpose cannot be moved before a tiling transform when it
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Unsupported transform: {type(transform)}
• Unsupported barrier type: {type(barrier)}

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