LogicalResult ResourceLiftingForFunctionalControlFlow(func::FuncOp function) {
  // This routine should only be called when control flow operations are still
  // represented with TF IfOp and WhileOp operations. In this case, there should
  // be only one basic blocks in the MLIR representation.
  if (!llvm::hasSingleElement(function)) {
    return function.emitError()
           << "expect the function to have 1 block while it has "

  // we do not overwrite the function argument mapping here.
  func_op.walk([&](Operation* op) {
    if (op == func_op) return;
    for (Region& region : op->getRegions()) {
      assign_unique_id_to_all(region.getArguments());
    }
  });

  llvm::SmallDenseMap<ResourceHandle, int64_t> resource_handle_id_map;
  func_op.walk([&](Operation* op) {

// tf_device.replicate.
Operation* FindCompilePredecessor(Operation* func_op,
                                  Operation* preprocess_op) {
  llvm::DenseSet<Block*> blocks = GetAllBlocksBetween(preprocess_op, func_op);

  llvm::DenseMap<Block*, Operation*> scope;
  Operation* o = preprocess_op;
  while (o && o != func_op) {
    scope[o->getBlock()] = o;
    o = o->getParentOp();
  }

  Value input = func_op.getArgument(0);
  Value output_init_state = func_op.getArgument(1);
  Value hidden_init_state = func_op.getArgument(2);
  Value weight_kernel = func_op.getArgument(3);
  Value recurrent_kernel = func_op.getArgument(4);
  Value bias = func_op.getArgument(5);

  // The func op should have 5 outputs.
  if (func_op.getNumResults() != 5) return failure();

  return IsReplicatedGraph(module);
}

bool HasTPUPartitionedCallOpInModule(mlir::ModuleOp module) {
  bool has_tpu_partitioned_call = false;
  for (auto func_op : module.getOps<mlir::func::FuncOp>()) {
    func_op->walk([&](mlir::TF::TPUPartitionedCallOp op) {
      has_tpu_partitioned_call = true;
    });
    if (has_tpu_partitioned_call) break;
  }
  return has_tpu_partitioned_call;
}

    return LogicalResult::success(success);
  }
};

class HandleFunc : public OpRewritePattern<func::FuncOp> {
 public:
  using OpRewritePattern::OpRewritePattern;

  LogicalResult matchAndRewrite(func::FuncOp f,
                                PatternRewriter& rewriter) const override {
    // Adjust the function type to match the block args. So this e.g. transforms
    //
    // func.func @f(%x : foo) -> foo {

                          std::unique_ptr<Graph>* graph,
                          FunctionLibraryDefinition* flib_def);

// Converts an MLIR function and adds it to a FunctionLibraryDefinition.
Status ConvertMlirFunctionToFunctionLibraryDef(mlir::func::FuncOp func,
                                               const GraphExportConfig& configs,
                                               FunctionDef* function_def);

  let constructor = "TFDevice::CreateResourceOpLiftingPass()";
}

def ResourceOpLiftingForMainFunctionPass :
    Pass<"tf-resource-op-lifting-for-main-function", "ModuleOp"> {
  let summary = "Lifting resource operations out of control flow statements "
    "for the main function";
  let constructor = "TFDevice::CreateResourceOpLiftingForMainFunctionPass()";
}

def AnnotateParameterReplicationPass :

// returned value are function name -> alias mappings. `function_aliases` is
// the function alias mapping of the original function. The original function's
// name is retrieved by looking at the "tf._original_func_name" string attribute
// attached to a `func::FuncOp`.
void UpdateFunctionAliases(
    absl::flat_hash_map<FunctionName, FunctionAlias>& function_aliases,
    ModuleOp module_op);

  // Optimizes TF graph via cleanups, merges, rewrites, constant folding,
  // and edge case handling where possible.
  pm.addNestedPass<func::FuncOp>(TF::CreateDropWhileShapeInvariantPass());
  pm.addNestedPass<func::FuncOp>(
      tf_executor::CreateTFExecutorGraphPruningPass());
  pm.addNestedPass<func::FuncOp>(
      tf_executor::CreateTFExecutorIslandCoarseningPass());
  pm.addPass(TF::CreateTFFunctionalControlFlowToRegions());