auto f_attr = mlir::dyn_cast<FlatSymbolRefAttr>(call_op.getFAttr());
    auto module_op = call_op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module_op);

    auto f_name = f_attr.getValue();
    func::FuncOp float_func =
        dyn_cast<func::FuncOp>(symbol_table.lookup(f_name));
    if (!float_func) {
      return failure();
    }
    rewriter.setInsertionPointAfter(call_op);

  void runOnOperation() override;

 private:
  LogicalResult ConvertIfOp(SymbolTableCollection& symbol_table,
                            IfRegionOp if_region);
  LogicalResult ConvertCaseOp(SymbolTableCollection& symbol_table,
                              CaseRegionOp case_region);
  LogicalResult ConvertWhileOp(SymbolTableCollection& symbol_table,
                               WhileRegionOp while_region);

void ColocateTPUCopyWithDynamicShapePass::runOnOperation() {
  Operation *module = getOperation();

  SymbolTableCollection symbolTables;

  DataFlowSolver solver;
  solver.load<dataflow::DeadCodeAnalysis>();
  solver.load<dataflow::SparseConstantPropagation>();
  solver.load<DeviceDataflowAnalysis>(symbolTables);
  if (failed(solver.initializeAndRun(module))) return signalPassFailure();

  while (!func_ops.empty()) {
    auto main_func = func_ops.back();
    func_ops.pop_back();
    if (!main_func) continue;

    SymbolTable symbol_table(module_op);
    for (auto call_op : main_func.getOps<TF::PartitionedCallOp>()) {
      func_ops.push_back(dyn_cast_or_null<func::FuncOp>(symbol_table.lookup(
          mlir::cast<FlatSymbolRefAttr>(call_op.getFAttr()).getValue())));
    }

  module_for_func.get().getOperation()->setAttr(kVersionsAttr, versions_attr);
  SymbolTable symbol_table(module_for_func.get());

  while (!referenced.empty()) {
    auto func = referenced.pop_back_val();

    // Skip functions that have already been cloned into new module.
    if (symbol_table.lookup<func::FuncOp>(func.getName())) continue;

  // bound_inputs' types correctly.
  SymbolTable symbol_table(module);
  for (auto func : module.getOps<func::FuncOp>()) {
    for (auto arg : func.getArguments()) {
      unsigned arg_number = arg.getArgNumber();
      auto global_tensor = LookupBoundInputOfType<GlobalTensorOp>(
          func, arg_number, symbol_table);
      if (!global_tensor) continue;

                               const StringRef func_name) {
  static tensorflow::mutex* mtx = new tensorflow::mutex();
  tensorflow::mutex_lock lock(*mtx);

  SymbolTable symbol_table(&module);
  std::string unique_name = func_name.str();
  int32_t uniquing_counter = 0;
  while (symbol_table.lookup(unique_name) != nullptr) {
    ++uniquing_counter;
    unique_name = absl::StrCat(func_name.str(), "_", uniquing_counter);
  }

                              StringRef error_message) const {
    ModuleOp module = call_op->getParentOfType<ModuleOp>();
    SymbolTable symbol_table(module);
    mlir::func::FuncOp composite_func =
        dyn_cast<func::FuncOp>(symbol_table.lookup(function_name));
    if (!composite_func) return;

    composite_func.walk([&](Operation* op) {
      if (op->hasAttr(kAttrMapAttribute)) {

      builder.getStringAttr(kTfSavedModelInitializerRestoreType));

  builder.setInsertionPointToStart(&init_func.front());
  builder.create<func::ReturnOp>(loc, /*operands=*/ValueRange{});

  SymbolTable symbol_table(module_op);
  symbol_table.insert(init_func);

  AddSymbolToInitializersAttr(
      session_init_op, FlatSymbolRefAttr::get(init_func.getSymNameAttr()));

  return init_func;
}

 private:
  void runOnOperation() override;
};

void ConvertXlaCallModuleOpToBfloat16Pass::runOnOperation() {
  Operation* func_op = getOperation();
  SymbolTableCollection symbol_table;
  OpBuilder builder(&getContext());

  auto result = func_op->walk([&](TF::XlaCallModuleOp op) {
    // Converts the serialized StableHLO module to bfloat16.
    auto result =