void runOnOperation() override;
};

void ResourceOpLiftingForMainFunctionPass::runOnOperation() {
  ModuleOp module = getOperation();
  func::FuncOp main_func = module.lookupSymbol<func::FuncOp>("main");
  if (!main_func) {
    return;
  }

  if (failed(TF::ResourceLiftingForFunctionalControlFlow(main_func))) {
    return signalPassFailure();
  }
}

    if (called_computations.size() != 1) {
      return op.emitOpError()
             << "ApproxTopK takes exactly 1 called_computation.";
    }
    mlir::func::FuncOp callee = module_op_->lookupSymbol<mlir::func::FuncOp>(
        mlir::cast<FlatSymbolRefAttr>(op.getCalledComputations()[0]));
    mlir::FunctionType callee_type = callee.getFunctionType();
    SmallVector<Type, 4> expected_callee_input_types;