// Move all remaining nested functions back into the parent module.
  Block &nested_block = nested_module->getRegion(0).front();
  for (func::FuncOp func_op :
       llvm::make_early_inc_range(nested_block.getOps<func::FuncOp>())) {
    if (!symbol_table.lookupSymbolIn(getOperation(), func_op.getName())) {
      nested_block.getOperations().remove(func_op.getOperation());
      symbol_table.insert(func_op.getOperation());
    }
  }

      init_func_op.emitError("expects exactly one block in the MLIR function");
      return signalPassFailure();
    }

    auto var_handle_ops =
        init_func_op.getBlocks().front().getOps<VarHandleOp>();
    llvm::SmallVector<VarHandleOp, 4> init_vars(var_handle_ops.begin(),
                                                var_handle_ops.end());
    RemoveVariables(init_vars);
  }
}

}  // namespace

      }
    }
  }
}

void AddFunctionsForExportedNamesPass::runOnOperation() {
  Block& module_body = getOperation().getRegion().front();
  for (auto f :
       llvm::make_early_inc_range(getOperation().getOps<func::FuncOp>())) {
    auto exported_names = GetExportedNames(f);
    if (exported_names.empty() || !f.isPublic()) continue;

    if (exported_names.size() == 1 && exported_names[0] == f.getName()) {

  if (!bundle_writer.status().ok()) {
    return bundle_writer.status();
  }

  std::vector<std::string> saved_variable_shared_names;
  for (auto assign_variable_op :
       session_init_func_type_restore_op.getOps<mlir::TF::AssignVariableOp>()) {
    if (const absl::StatusOr<std::string> variable_shared_name =
            AddTensorToBundleWriter(assign_variable_op, bundle_writer);
        !variable_shared_name.ok()) {

      return failure();

    do {
      SymbolUserMap symbol_users(symbol_table_collection, module);

      made_changes = false;
      for (auto func :
           llvm::make_early_inc_range(module.getOps<func::FuncOp>())) {
        ArrayRef<Operation *> users = symbol_users.getUsers(func);
        if (users.size() <= 1) {
          continue;
        }

  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LegalizeWhilePass)
  void RunOnFunction(func::FuncOp func);

  void runOnOperation() override {
    for (auto op : getOperation().getOps<func::FuncOp>()) RunOnFunction(op);
  }
};

}  // namespace

// Inserts call to the given function into the 'region'.
void CreateRegionWithCall(func::FuncOp func, Region& region, Location loc) {

        device_mgr_.get(), Env::Default(), /*config=*/nullptr,
        TF_GRAPH_DEF_VERSION, lib_def_.get(), opts,
        /*default_thread_pool=*/nullptr, /*cluster_flr=*/nullptr);
    flr_ = pflr_->GetFLR("/job:localhost/replica:0/task:0/cpu:0");
  }

  FunctionLibraryRuntime* flr_;
  std::unique_ptr<DeviceMgr> device_mgr_;
  std::unique_ptr<FunctionLibraryDefinition> lib_def_;

// failure.
template <typename OpT>
LogicalResult GetOp(Region* region, OpT* result) {
  *result = {};
  for (auto op : region->getOps<OpT>()) {
    if (*result) return op.emitError("should be unique within a function");
    *result = op;
  }
  return success();
}

LogicalResult RunOnRegion(Region* region) {
  RecvTPUEmbeddingActivationsOp recv_op;
  if (failed(GetOp(region, &recv_op))) return failure();

      }))
    return replicated_input.emitOpError()
           << "expects all inputs from 'tf.TPUPartitionedInputV2' ops";

  const auto metadata_iter =
      replicated_input->getBlock()->getOps<TF::TPUReplicateMetadataOp>();
  TF::TPUReplicateMetadataOp metadata;
  if (!metadata_iter.empty()) metadata = *(metadata_iter.begin());

  auto first_partitioned_input = llvm::cast<TF::TPUPartitionedInputV2Op>(

    EXPECT_TRUE(module_op_ref);
    return module_op_ref;
  }

  mlir::MLIRContext ctx_{};
};

func::FuncOp GetMainFuncOp(ModuleOp module_op) {
  for (auto func_op : module_op.getOps<func::FuncOp>()) {
    if (func_op.getSymName() == "main") {
      return func_op;
    }
  }
  return {};
}

TEST_F(ConvertAssetArgsTest, ConvertsSingleAssetArg) {