jax-ml/jax · error · ValueError

Expected a dynamic index or an integer, got {offset}

Error message

Expected a dynamic index or an integer, got {offset}

What it means

Barrier offsets must be Python ints, index-typed values, or i32 values. Anything else (i64, f32, tensor, etc.) reaches the final else and raises ValueError telling you what type was passed.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:1077

    return BarrierRef(address, c(0, i32), phases, num_barriers)

  def __iter__(self) -> Iterator["BarrierRef"]:
    if self.num_barriers == 1:
      yield self
    else:
      for offset in range(self.num_barriers):
        yield self[offset]

  def __getitem__(self, offset: ir.Value | int) -> "BarrierRef":
    i32 = ir.IntegerType.get_signless(32)
    if isinstance(offset, int):
      if offset >= self.num_barriers:
        raise IndexError(f"Barrier offset {offset} is out of bounds")
      offset = c(offset, i32)
    elif isinstance(offset.type, ir.IndexType):
      offset = arith.index_castui(i32, offset)
    elif offset.type != i32:
      raise ValueError(f"Expected a dynamic index or an integer, got {offset}")
    return BarrierRef(
        self.base_address,
        arith.addi(self.offset, offset),
        self.phases,
        1,
    )

  @property
  def _ptx_scope(self) -> str:
    if self.base_address.type == ir.Type.parse("!llvm.ptr<7>"):
      return "cluster"
    return "cta"

  @property
  def _nvvm_scope(self) -> nvvm.MemScopeKind:
    if self.base_address.type == ir.Type.parse("!llvm.ptr<7>"):
      return nvvm.MemScopeKind.CLUSTER
    return nvvm.MemScopeKind.CTA

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast to i32: arith.trunci(ir.IntegerType.get_signless(32), val)
2. Use index type for dynamic offsets (auto-cast via index_castui) or plain Python ints for static ones
3. Validate with isinstance(offset, int) or offset.type in (i32, IndexType) before indexing

Example fix

# before
ref = group[offset_i64]
# after
offset = arith.trunci(ir.IntegerType.get_signless(32), offset_i64)
ref = group[offset]

i32 = ir.IntegerType.get_signless(32)
if isinstance(offset, ir.Value) and offset.type not in (i32, ir.IndexType.get()):
    offset = arith.trunci(i32, offset)

def valid_barrier_offset(o) -> bool:
    i32 = ir.IntegerType.get_signless(32)
    return isinstance(o, int) or (isinstance(o, ir.Value) and o.type in (i32, ir.IndexType.get()))

When it happens

Trigger: group[i64_value] where i64_value = arith.constant(1, i64); group[1.0]; passing an unconverted OpResult of non-i32 type.

Common situations: Computing barrier offsets with 64-bit arithmetic (e.g. byte offsets or grid math) and forgetting to cast; mixing pointer/integer width conventions between PTX asm helpers and the barrier API.

Related errors

• Indexing barrier with {transforms} not supported in GPU inte
• Expected an `NDIndexer`, but got {transforms[0]}
• Attempting to operate on barrier without indexing, but `num_
• Expected a single barrier, got a barrier reference with shap
• copy_gmem_to_smem with a barrier is only supported Hopper an

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