func::FuncOp cond =
      func::FuncOp::create(loc, "new_while_cond", cond_func_type);
  func::FuncOp body =
      func::FuncOp::create(loc, "new_while_body", body_func_type);
  cond.setPrivate();
  body.setPrivate();
  symbol_table.insert(cond);
  symbol_table.insert(body);
  OpBuilder cond_builder = OpBuilder::atBlockBegin(cond.addEntryBlock());

    const InferenceDeviceType& target_device_inference_type, ModuleOp module,
    OpBuilder* builder) {
  func::FuncOp cloned_func = func.clone();
  cloned_func.setPrivate();
  auto interface_name = GetInterFaceName(func);
  if (!interface_name.has_value()) {
    func.emitError("the func op does not have interface_name");
    return nullptr;
  }

  cloned_func->setAttr(

  auto func_uses = SymbolTable::getSymbolUses(func, &module.getBodyRegion());
  if (func_uses.has_value() && llvm::hasSingleElement(func_uses.value()))
    return func;
  func::FuncOp cloned = func.clone();
  cloned.setPrivate();
  cloned.setName(
      StringAttr::get(func.getContext(), func.getName().str() + "_lifted"));
  SymbolTable(module).insert(cloned);
  return cloned;
}

      // should be no other reference to it.
      clone.setName(StringAttr::get(clone.getContext(), "main"));
      clone.setPublic();
    } else {
      clone.setPrivate();
    }
    symbol_table.insert(clone);
  }

  *serialized_func_module =
      tensorflow::SerializeMlirModule(module_for_func.get());
  return success();
}

  func::FuncOp lowered_callee = callee;
  if (!callee.isPrivate()) {
    // Clone non-private callee in case of signature change.
    lowered_callee = callee.clone();
    lowered_callee.setPrivate();
  }
  auto find_arg_stack_type = [&](int64_t index) -> std::optional<Type> {
    auto it = data_var_to_size_var.find(call.getOperand(index));
    if (it == data_var_to_size_var.end()) return std::nullopt;

  terminator = first_block.getTerminator();
  builder.setInsertionPoint(terminator);
  builder.create<func::ReturnOp>(terminator->getLoc(), return_values);
  terminator->erase();

  outlined_func.setPrivate();

  // Uniquify the function name, and insert into module.
  symbol_table.getSymbolTable(module).insert(outlined_func,
                                             module.getBody()->begin());

  func::FuncOp lowered_callee = callee;
  if (!callee.isPrivate()) {
    // Clone non-private callee in case of signature change.
    lowered_callee = callee.clone();
    lowered_callee.setPrivate();
  }
  auto find_arg_buffer_type = [&](int64_t index) -> std::optional<Type> {
    auto it = buffer_to_size->find(call.getOperand(index));
    if (it == buffer_to_size->end()) return std::nullopt;

  builder.setInsertionPointToEnd(&module.getBodyRegion().back());
  func::FuncOp func_op = builder.create<func::FuncOp>(
      module.getLoc(), name,
      builder.getFunctionType(input_types, output_types));
  func_op.setPrivate();

  return func_op;
}

TF::StatefulPartitionedCallOp EncapsulateOpsInFunc(
    OpBuilder& builder, const llvm::SetVector<Operation*>& ops,

  for (auto& func : funcs_to_rewrite) {
    if (func.getSecond().original.isPublic() &&
        !func.getSecond().original.symbolKnownUseEmpty(module)) {
      auto clone = func.getSecond().original.clone();
      clone.setPrivate();
      symbol_table.insert(clone);
      func.getSecond().clone = clone;
    }
  }

  return success();
}

// Assigns either MAXIMAL or MANUAL sharding. The MAXIMAL sharding sends/recvs

  };
  func::FuncOp lowered_callee = callee;
  if (!callee.isPrivate()) {
    // Clone non-private callee in case of signature change.
    lowered_callee = callee.clone();
    lowered_callee.setPrivate();
  }
  auto grads = AccessedGradients({lowered_callee}, module);
  for (int64_t i = 0; i < lowered_callee.getNumArguments(); ++i) {
    auto it = grads.find(i);
    if (it == grads.end()) continue;