{"record":{"id":"d0d2620b654692c1","repo":"jax-ml/jax","slug":"smem-has-shape-smem-shape-but-expected-expecte-d0d262","errorCode":null,"errorMessage":"SMEM has shape {smem_shape}, but expected {expected_smem_shape} for TMEM shape {tmem_ref.shape} with swizzle={swizzle}","messagePattern":"SMEM has shape (.+?), but expected (.+?) for TMEM shape (.+?) with swizzle=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":2025,"sourceCode":"  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\n\n  # Figure out the matrix descriptor parameters (LBO/SBO)\n  # The copy happens using the usual \"core matrix\" structure: a memory region\n  # describing a 8x128bit matrix. LBO describes how far apart from each other\n  # are consecutive matrices along the minor dimension (in our case the minor\n  # dim is contiguous, so exactly 128 bit = 16 bytes apart). SBO describes how\n  # far apart is the beginning of the next matrix along the major dimension.","sourceCodeStart":2007,"sourceCodeEnd":2043,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L2007-L2043","documentation":"Raised by async_copy_smem_to_tmem when the SMEM memref's shape does not equal the tile shape implied by the TMEM reference and swizzle: expected_smem_shape = tile_shape(tmem_ref.shape, (8, swizzle_elems)) where swizzle_elems = 8*swizzle/bitwidth. The instruction walks SMEM as 8x(swizzle_elems) core matrices, so the buffer must be exactly that tiled shape.","triggerScenarios":"Allocating SMEM as a flat (rows, cols) array instead of the 4D tiled shape (row_tiles, col_tiles, 8, swizzle_elems), or passing a swizzle whose byte size mismatches the SMEM tile width, e.g. smem shape (128, 64) with swizzle=32 and f32 elements (swizzle_elems=8, so tiles are 8x8 and total shape must reflect that).","commonSituations":"Porting wgmma/legacy mma SMEM staging layouts to tcgen05; changing element dtype without re-deriving swizzle_elems (swizzle is in bytes, shape is in elements).","solutions":["Allocate SMEM with shape = tile_shape(tmem_shape, (8, 8*swizzle//bitwidth(dtype))) — i.e. reshape your data into 8x(swizzle_elems) core-matrix tiles","Recompute the expected shape in Python with utils.tile_shape and assert your smem allocation matches before lowering","If dtype changes, keep swizzle fixed and let the tile shape change accordingly"],"exampleFix":"# before\nswizzle = 128; dtype_bits = 32\nsmem = smem_alloc(f32, (128, 64))\n# after\nfrom jax.experimental.mosaic.gpu import utils\nse = 8 * swizzle // dtype_bits\nsmem_shape = utils.tile_shape(tmem_ref.shape, (8, se))\nsmem = smem_alloc(f32, smem_shape)","handlingStrategy":"validation","validationCode":"bw = utils.bitwidth(dtype)\nexpected = utils.tile_shape(tmem_ref.shape, (8, 8 * swizzle // bw))\nassert tuple(ir.MemRefType(smem_ref.type).shape) == expected, f'smem shape must be {expected}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive the SMEM tiled shape from the TMEM shape and swizzle with utils.tile_shape, never hand-write it","When dtype changes, recompute swizzle_elems since swizzle is bytes but shape is elements"],"tags":["gpu","mosaic","tcgen05","shape-mismatch","swizzle","shared-memory"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}