{"record":{"id":"0e7aec2d19901b12","repo":"jax-ml/jax","slug":"only-standard-tmem-layout-is-supported-got-tmem","errorCode":null,"errorMessage":"Only standard TMEM layout is supported, got: {tmem_ref.layout}","messagePattern":"Only standard TMEM layout is supported, got: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":2018,"sourceCode":"  i8 = ir.IntegerType.get_signless(8)\n  i32 = ir.IntegerType.get_signless(32)\n  smem_ty = ir.MemRefType(smem_ref.type)\n  if (dtype := smem_ty.element_type) != tmem_ref.dtype:\n    raise ValueError(f\"Incompatible dtypes: SMEM has {dtype}, TMEM has {tmem_ref.dtype}\")\n  if swizzle not in {16, 32, 64, 128}:\n    raise ValueError(f\"Unsupported swizzle, expected 16, 32, 64 or 128, but got: {swizzle}\")\n  bitwidth = utils.bitwidth(dtype)\n  if tmem_ref.packing != 32 // bitwidth:\n    raise ValueError(\n        \"tcgen05.cp only supports fully packed TMEM references\"\n        f\" (packing={32 // bitwidth}), but got packing={tmem_ref.packing}\"\n    )\n  if tmem_ref.shape[0] != TMEM_ROWS:\n    raise ValueError(\n        f\"TMEM reference must have {TMEM_ROWS} rows, but got {tmem_ref.shape[0]}\"\n    )\n  if tmem_ref.layout != tmem_default_layout(packing=tmem_ref.packing):\n    raise ValueError(\n        f\"Only standard TMEM layout is supported, got: {tmem_ref.layout}\"\n    )\n  swizzle_elems = 8 * swizzle // bitwidth\n  expected_smem_shape = utils.tile_shape(tmem_ref.shape, (8, swizzle_elems))\n  smem_shape = tuple(smem_ty.shape)\n  if smem_shape != expected_smem_shape:\n    raise ValueError(\n        f\"SMEM has shape {smem_shape}, but expected {expected_smem_shape} for\"\n        f\" TMEM shape {tmem_ref.shape} with swizzle={swizzle}\"\n    )\n  strides, _ = smem_ty.get_strides_and_offset()\n  row_tile_stride, col_tile_stride, inner_row_stride, inner_col_stride = strides\n  if inner_col_stride != 1 or inner_row_stride != swizzle_elems:\n    raise ValueError(\"The SMEM tiles must be contiguous\")\n  # Make sure strides are a multiple of the byte packing for narrow types.\n  byte_packing = max(8 // bitwidth, 1)\n  assert row_tile_stride % byte_packing == 0\n  assert col_tile_stride % byte_packing == 0","sourceCodeStart":2000,"sourceCodeEnd":2036,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L2000-L2036","documentation":"Raised by async_copy_smem_to_tmem when the TMEM reference's layout does not match tmem_default_layout(packing=...). tcgen05.cp only emits correct addresses for the standard (default) TMEM layout; custom blocked/interleaved TMEM layouts are not supported for this copy instruction.","triggerScenarios":"Passing a tmem_ref created with a custom layout argument to tmem_alloc (or one produced by tmem slice ops carrying a non-default layout) into async_copy_smem_to_tmem.","commonSituations":"Experimenting with TMEM layouts for accumulator reuse across MMAs, then feeding the same reference into a tcgen05.cp; defaulting to a legacy layout after a JAX/Mosaic version change altered default layouts.","solutions":["Omit the layout argument in tmem_alloc so the default layout is used, or pass tmem_default_layout(packing=32//bitwidth(dtype)) explicitly","Copy into a fresh default-layout TMEM reference and move data with TMEM-to-TMEM or register ops if a custom layout is required downstream"],"exampleFix":"// before\ntmem = tmem_alloc(dtype, shape, layout=custom_layout)\n// after\nfrom jax.experimental.mosaic.gpu import tcgen05\ntmem = tmem_alloc(dtype, shape, layout=tcgen05.tmem_default_layout(packing=32//bitwidth(dtype)))","handlingStrategy":"validation","validationCode":"expected = tcgen05.tmem_default_layout(packing=tmem_ref.packing)\nassert tmem_ref.layout == expected, 'non-default TMEM layout cannot be a tcgen05.cp destination'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Don't attach custom layouts to TMEM refs destined for async_copy_smem_to_tmem","Keep a dedicated default-layout TMEM buffer for SMEM->TMEM staging"],"tags":["gpu","mosaic","tcgen05","layout","tensor-memory"],"backgroundTag":"invalid-layout-configuration","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}