{"record":{"id":"e8fa64c99ee79381","repo":"jax-ml/jax","slug":"expected-tiledlayout-got-type-layout","errorCode":null,"errorMessage":"Expected TiledLayout, got {type(layout)}","messagePattern":"Expected TiledLayout, got (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":601,"sourceCode":"    raise ValueError(f\"Unsupported memory space: {ref_type.memory_space}\")\n\n  if ctx.auto_barriers:\n    utils.warpgroup_barrier()  # Make sure the writes have completed.\n\n  return []\n\n\n@_register_lowering(mgpu.AsyncStoreSmemOp)\ndef _async_store_smem_op_lowering_rule(\n    ctx: LoweringContext, op: mgpu.AsyncStoreSmemOp\n) -> Sequence[ir.Value]:\n  index = ir.IndexType.get()\n\n  [to_store_layout] = inference_utils.in_layouts(op)\n  value = _fragmented_array_from_ir(op.valueToStore, to_store_layout)\n  layout = layouts_lib.from_layout_attr(to_store_layout)\n  if not isinstance(layout, fa.TiledLayout):\n    raise NotImplementedError(f\"Expected TiledLayout, got {type(layout)}\")\n\n  ref = op.destination\n  transforms_attr = inference_utils.in_transforms(op)[0]\n  swizzle = swizzle_from_transforms_attr(transforms_attr)\n  unwrapped_ref = unwrap_transformed_memref(ref, transforms_attr)\n  tiling_transform, = memref_transforms_from_transforms_attr(transforms_attr)\n  assert isinstance(tiling_transform, lc.TileTransform)\n\n  dialect_barrier = utils.DialectBarrierRef.from_barrier_memref(op.barrier)\n  barrier_ref = dialect_barrier.barrier_ref\n\n  cluster_dim = gpu.Dimension(op.cluster_dim.value)  # pyrefly: ignore[missing-attribute]\n  cluster_idx = arith.index_cast(index, op.cluster_idx)\n  cluster_barrier_ref = barrier_ref.remap_to_cluster(cluster_dim, cluster_idx)\n\n  total_bits = math.prod(value.shape) * utils.bitwidth(value.mlir_dtype)\n  if total_bits % (8 * utils.WARPGROUP_SIZE):\n    raise NotImplementedError(","sourceCodeStart":583,"sourceCodeEnd":619,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L583-L619","documentation":"mgpu.async_store_smem requires the value's layout to be a TiledLayout; fragment-based layouts cannot be written to shared memory by this op.","triggerScenarios":"Calling async_store_smem (or emitting mgpu.async_store_smem) where the in_layout attr decodes to a fragment layout like WGStridedFragLayout.","commonSituations":"Storing a matmul result (accumulator fragment layout) directly to SMEM without first casting to a tiled layout.","solutions":["Cast the value to a TiledLayout (layout_cast / to_layout) before async_store_smem","Check isinstance(layouts_lib.from_layout_attr(layout), TiledLayout) before emitting the op"],"exampleFix":"// before\nasync_store_smem(value, smem_ref)\n// after\ntiled = layout_cast(value, tiled_layout)\nasync_store_smem(tiled, smem_ref)","handlingStrategy":"type-guard","validationCode":"layout = layouts_lib.from_layout_attr(in_layout)\nassert isinstance(layout, fa.TiledLayout), 'async_store_smem requires TiledLayout'","typeGuard":"def is_tiled(attr) -> bool:\n    return isinstance(layouts_lib.from_layout_attr(attr), fa.TiledLayout)","tryCatchPattern":null,"preventionTips":["layout_cast fragment values to tiled before async_store_smem"],"tags":["jax","mosaic","gpu","smem","layout"],"backgroundTag":"unsupported-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}