{"record":{"id":"b85a937ed788e085","repo":"jax-ml/jax","slug":"collective-argument-must-be-provided-when-tmem-lay","errorCode":null,"errorMessage":"collective argument must be provided when TMEM layout is inferred","messagePattern":"collective argument must be provided when TMEM layout is inferred","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1234,"sourceCode":"      layout: TMEMLayout | None = None,\n  ) -> TMEMRef:\n    i32 = ir.IntegerType.get_signless(32)\n    if not isinstance(tmem_addr_ref.type, ir.MemRefType):\n      raise ValueError(f\"tmem_addr_ref must be a memref or a pointer, got: {tmem_addr_ref.type}\")\n    addr_ref_ty = ir.MemRefType(tmem_addr_ref.type)\n    if not utils.is_smem_ref(addr_ref_ty):\n      raise ValueError(f\"tmem_addr_ref must be in shared memory, got: {addr_ref_ty}\")\n    if addr_ref_ty.element_type != i32:\n      raise ValueError(f\"tmem_addr_ref must be an i32 memref, got: {addr_ref_ty}\")\n    if math.prod(addr_ref_ty.shape) != 1:\n      raise ValueError(f\"tmem_addr_ref must contain a single element, got: {addr_ref_ty}\")\n    i0 = arith.ConstantOp.create_index(0)\n    tmem_addr = memref.load(tmem_addr_ref, [i0] * addr_ref_ty.rank)\n    if shape[0] < 32:\n      raise ValueError(f\"TMEM refs must have at least 32 rows, got: {shape[0]}\")\n    if layout is None:\n      if collective is None:\n        raise ValueError(\n            \"collective argument must be provided when TMEM layout is inferred\"\n        )\n      layout = _infer_tmem_layout(shape, collective, packing=1)\n    # TODO: Do we have to do this??\n    # warp_idx = utils.warp_idx(sync=False)\n    # tmem_addr = arith.ori(tmem_addr, arith.shli(warp_idx, utils.c(21, i32)))\n    return cls(tmem_addr, shape, dtype, layout)\n\n  def slice(self, *idxs) -> TMEMRef:\n    i32 = ir.IntegerType.get_signless(32)\n    base_idx, slice_shape, is_squeezed = utils.parse_indices(idxs, self.shape)\n    slice_shape = cast(tuple[int, int], tuple(slice_shape))\n    if any(is_squeezed):\n      raise ValueError(\"TMEM can only be sliced, not indexed\")\n    if base_idx == [0] * len(base_idx) and slice_shape == self.shape:\n      return self  # Trivial slice\n    # If we slice along rows, or attempt to extract several rows, then we may\n    # end up with a non-contiguous slice of memory.","sourceCodeStart":1216,"sourceCodeEnd":1252,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1216-L1252","documentation":"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.","triggerScenarios":"TMEMRef.from_alloc(alloc, shape) with neither layout= nor collective= given.","commonSituations":"Simplified examples omitting collective; refactoring that dropped the collective argument when moving from explicit layouts to inferred ones.","solutions":["Pass collective=True/False explicitly when omitting layout","Or pass an explicit layout= (e.g. tmem_default_layout()) so no inference is needed"],"exampleFix":"# before\nref = tcgen05.TMEMRef.from_alloc(alloc, (64, 64))\n# after\nref = tcgen05.TMEMRef.from_alloc(alloc, (64, 64), collective=False)","handlingStrategy":"validation","validationCode":"assert layout is not None or collective is not None","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always pass collective= (True for 2CTA MMA) unless you pass an explicit layout","Wrap from_alloc in a project helper that requires one of the two"],"tags":["jax","mosaic","tmem","missing-argument"],"backgroundTag":"missing-required-argument","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}