{"record":{"id":"204eec5ec4475c34","repo":"jax-ml/jax","slug":"expanding-tiled-dimensions-is-not-supported","errorCode":null,"errorMessage":"Expanding tiled dimensions is not supported.","messagePattern":"Expanding tiled dimensions is not supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2226,"sourceCode":"  in_transforms = inference_utils.in_transforms(op)[0]\n  unwrapped_in_ref = unwrap_transformed_memref(op.src, in_transforms)\n  in_transformed_ty = ir.MemRefType(unwrapped_in_ref.type)\n\n  out_transforms = inference_utils.out_transforms(op)[0]\n  out_transformed_ty = transform_type(op.result.type, out_transforms)\n\n  reassociation = cast(list[ir.ArrayAttr], list(op.reassociation))\n  num_tiling_dims = len(in_transformed_ty.shape) - len(op.src.type.shape)\n\n  # We don't currently allow expanding tiled dimensions. So to compute the\n  # reassociation on the lowered types, we just need to backfill the original\n  # one with the number of missing dimensions.\n  if num_tiling_dims > 0 and any(\n      len(x) > 1 for x in reassociation[-num_tiling_dims:]\n  ):\n    # If we ever remove this restriction, we will need to ensure this is\n    # compatible with `transform_type`.\n    raise NotImplementedError(\"Expanding tiled dimensions is not supported.\")\n\n  start_index = len(op.static_output_shape)\n  for i in range(start_index, start_index + num_tiling_dims):\n    reassociation.append([i])  # pyrefly: ignore[bad-argument-type]\n\n  new_expand_shape_op = memref.ExpandShapeOp(\n      out_transformed_ty,\n      unwrapped_in_ref,\n      reassociation,\n      output_shape=op.output_shape,\n      static_output_shape=out_transformed_ty.shape,\n  )\n\n  wrapped_ref = wrap_transformed_memref(\n      new_expand_shape_op.result, op.result.type, out_transforms\n  )\n  return [wrapped_ref]\n","sourceCodeStart":2208,"sourceCodeEnd":2244,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2208-L2244","documentation":"memref.expand_shape lowering cannot split a tiled dimension into multiple dimensions: if any of the last num_tiling_dims reassociation groups has more than one member, expansion of tiled dims is rejected.","triggerScenarios":"memref.expand_shape on a tiled memref where a reassociation group covering a tiled dimension expands it into 2+ output dims (len(group) > 1).","commonSituations":"Reshaping a tiled smem tensor to expose inner axes (e.g. splitting the last tiled dim into two) for element-wise access; not supported due to transform_type incompatibility noted in the source comment.","solutions":["Expand only untiled dimensions; keep tiled dims intact (reassociation entry [i] alone)","Retile with smaller tiles before expanding so expansion targets untiled dims","Do the expansion before applying the tile transform"],"exampleFix":"// before\n# tiled last dim of size 8\nout = t.memref.expand_shape(tiled, reassociation=[[0],[1,2]])  # splits tiled dim\n// after\nout = t.memref.expand_shape(tiled, reassociation=[[0,1],[2]])  # split untiled dims only","handlingStrategy":"validation","validationCode":"assert all(len(g) == 1 for g in reassociation[-num_tiling_dims:]), 'cannot expand tiled dims'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Only expand untiled dimensions","Tile after expanding"],"tags":["jax","mosaic-gpu","expand-shape","tiling","not-implemented"],"backgroundTag":"unsupported-layout-transform","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}