ModuleOp module_op) {
  MLIRContext* ctx = module_op.getContext();
  OpBuilder builder(ctx);
  Operation* last_subgraph_op = reverse_subgraph.front();
  builder.setInsertionPointAfter(last_subgraph_op);

  // Create attributes used for creating an XlaCallModuleOp.
  SmallVector<Attribute> shape_attrs;
  for (const Type result_type : result_types) {
    shape_attrs.push_back(

      op->setOperand(
          arg, builder.create<TransposeOp>(loc, op->getOperand(arg), arg_perm));
    }

    // Permute results back to the original data format.
    builder.setInsertionPointAfter(op);
    auto res_perm = builder.create<ConstOp>(loc, perm_attr(res_permutation));

    for (int64_t res : layout_sensitive_interface.GetLayoutDependentResults()) {
      OpResult result = op->getResult(res);

    axis_stats = DenseFPElementsAttr::get(
        RankedTensorType::get(axis_stats_shape, b.getF32Type()), min_maxs);
    axis = b.getI64IntegerAttr(info.meta().quantize_axis());
  }

  b.setInsertionPointAfter(op);
  if (IsQuantizableResult(op, index)) {
    InsertStatsOpAtResult(b, op->getResult(index), layer_stats, axis_stats,
                          axis);
  } else {

          output_types.push_back(result.getType());
        } else {
          // TODO(b/224691264): separate matching and rewriting clearly.
          return failure();
        }
      }

      rewriter.setInsertionPointAfter(quantizing_op);
      OperationState new_state(quantizing_op->getLoc(),
                               quantizing_op->getName().getStringRef(), inputs,

  // Create a private function and copy all ops between arguments and results.
  auto current_func = result_op->getParentOfType<func::FuncOp>();
  auto guard = OpBuilder::InsertionGuard(builder);
  builder.setInsertionPointAfter(current_func);
  TypeRange arg_types{ValueRange{arguments}};
  TypeRange result_types{ValueRange{results}};
  auto func_type = FunctionType::get(context, arg_types, result_types);

    auto func = BuildGELUDecomposition(
        rhs_add.getType().cast<RankedTensorType>(), rewriter, root.getBody());

    SymbolTable table(root);
    (void)table.renameToUnique(func, {});

    rewriter.setInsertionPointAfter(output_mul);
    auto composite_attrs = rewriter.getDictionaryAttr(
        {rewriter.getNamedAttr("approx", rewriter.getBoolAttr(false))});
    auto composite_op = rewriter.create<stablehlo::CompositeOp>(

      return;
    }

    const Type expressed_type = op->getResult(0).getType();
    const Type quantized_type =
        quant_type.castFromExpressedType(expressed_type);

    rewriter.setInsertionPointAfter(op);
    auto q = rewriter.create<quantfork::QuantizeCastOp>(
        op->getLoc(), quantized_type, op->getResult(0));
    auto dq = rewriter.create<quantfork::DequantizeCastOp>(op->getLoc(),

  for (auto result : op->getResults()) {
    if (!erase[result.getResultNumber()]) {
      result_types.push_back(result.getType());
    }
  }
  OpBuilder builder(op);
  builder.setInsertionPointAfter(op);
  OperationState state(op->getLoc(), op->getName().getStringRef(),
                       op->getOperands(), result_types, op->getAttrs());
  for (int i = 0; i < op->getNumRegions(); ++i) {
    state.addRegion();

    auto shape_spec_type =
        RankedTensorType::get({cur_rank}, rewriter.getIntegerType(64));
    auto new_shape_const_attr =
        DenseElementsAttr::get(shape_spec_type, new_shape.getShape());
    rewriter.setInsertionPointAfter(weight_op);
    auto new_shape_const = rewriter.create<arith::ConstantOp>(
        weight_op->getLoc(), shape_spec_type, new_shape_const_attr);
    auto reshape_op = rewriter.create<TF::ReshapeOp>(

          outputs_replaced.insert({result, enumerated_result.index()});
          output_types.push_back(result.getType());
        } else {
          return failure();
        }
      }

      rewriter.setInsertionPointAfter(candidate_op);
      OperationState new_state(candidate_op->getLoc(),
                               candidate_op->getName().getStringRef(), inputs,
                               output_types, candidate_op->getAttrs());