jax-ml/jax · error · ValueError

collective argument must be provided when TMEM layout is inf

Error message

collective argument must be provided when TMEM layout is inferred

What it means

When layout is None, from_alloc infers the TMEM layout via _infer_tmem_layout, which must know whether the MMA is 2CTA (collective) to pick tmem_m64_collective_layout vs tmem_half_lane_layout for 64-row shapes. If you omit both layout and collective, inference is impossible.

Source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:1234

      layout: TMEMLayout | None = None,
  ) -> TMEMRef:
    i32 = ir.IntegerType.get_signless(32)
    if not isinstance(tmem_addr_ref.type, ir.MemRefType):
      raise ValueError(f"tmem_addr_ref must be a memref or a pointer, got: {tmem_addr_ref.type}")
    addr_ref_ty = ir.MemRefType(tmem_addr_ref.type)
    if not utils.is_smem_ref(addr_ref_ty):
      raise ValueError(f"tmem_addr_ref must be in shared memory, got: {addr_ref_ty}")
    if addr_ref_ty.element_type != i32:
      raise ValueError(f"tmem_addr_ref must be an i32 memref, got: {addr_ref_ty}")
    if math.prod(addr_ref_ty.shape) != 1:
      raise ValueError(f"tmem_addr_ref must contain a single element, got: {addr_ref_ty}")
    i0 = arith.ConstantOp.create_index(0)
    tmem_addr = memref.load(tmem_addr_ref, [i0] * addr_ref_ty.rank)
    if shape[0] < 32:
      raise ValueError(f"TMEM refs must have at least 32 rows, got: {shape[0]}")
    if layout is None:
      if collective is None:
        raise ValueError(
            "collective argument must be provided when TMEM layout is inferred"
        )
      layout = _infer_tmem_layout(shape, collective, packing=1)
    # TODO: Do we have to do this??
    # warp_idx = utils.warp_idx(sync=False)
    # tmem_addr = arith.ori(tmem_addr, arith.shli(warp_idx, utils.c(21, i32)))
    return cls(tmem_addr, shape, dtype, layout)

  def slice(self, *idxs) -> TMEMRef:
    i32 = ir.IntegerType.get_signless(32)
    base_idx, slice_shape, is_squeezed = utils.parse_indices(idxs, self.shape)
    slice_shape = cast(tuple[int, int], tuple(slice_shape))
    if any(is_squeezed):
      raise ValueError("TMEM can only be sliced, not indexed")
    if base_idx == [0] * len(base_idx) and slice_shape == self.shape:
      return self  # Trivial slice
    # If we slice along rows, or attempt to extract several rows, then we may
    # end up with a non-contiguous slice of memory.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass collective=True/False explicitly when omitting layout
2. Or pass an explicit layout= (e.g. tmem_default_layout()) so no inference is needed

Example fix

# before
ref = tcgen05.TMEMRef.from_alloc(alloc, (64, 64))
# after
ref = tcgen05.TMEMRef.from_alloc(alloc, (64, 64), collective=False)

assert layout is not None or collective is not None

When it happens

Trigger: TMEMRef.from_alloc(alloc, shape) with neither layout= nor collective= given.

Common situations: Simplified examples omitting collective; refactoring that dropped the collective argument when moving from explicit layouts to inferred ones.

Understand the failure class

Background: "missing required argument" and "the following required arguments were not provided": what required-argument errors mean and how to fix them — this error's family across 20 libraries.

Related errors

• Minor dimension of shape must be divisible by packing, got:
• Unsupported shape: {shape}. TMEM references must have either
• Packing must be a power of 2, got: {packing}
• Columns must be a multiple of 16, got: {columns}
• tmem_addr_ref must be a memref or a pointer, got: {tmem_addr

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