EXPECT_THAT(pre_calibration_result, IsOk());

  SmallVector<func::FuncOp> func_ops;
  for (auto func_op : pre_calibration_result->getOps<func::FuncOp>()) {
    func_ops.push_back(func_op);
  }
  ASSERT_THAT(func_ops, SizeIs(2));
  EXPECT_THAT(func_ops, Contains(HasSymName("main")));
  EXPECT_THAT(func_ops, Contains(HasSymName("composite_dot_general_fn_1")));

  func::FuncOp main_func = FindMainFuncOp(module_op);
  if (!main_func) return;

  // In case the model has tf.StatefulPartitionedCallOp or tf.PartitionedCallOp,
  // we recursively find called functions and process StableHLO ops in them.
  SmallVector<func::FuncOp> func_ops;
  func_ops.push_back(main_func);
  int stablehlo_func_id = -1;
  while (!func_ops.empty()) {
    auto main_func = func_ops.back();

};

// Returns true iff func_op has either no Region or the body has no Blocks.
bool IsFuncOpEmpty(func::FuncOp func_op) {
  return func_op->getNumRegions() == 0 || func_op.getBody().empty();
}

// Gets the GraphOp from the function op. Returns an empty op iff it doesn't
// exist.
GraphOp GetGraphOpFromFuncOp(func::FuncOp func_op) {
  if (IsFuncOpEmpty(func_op)) return {};

  return false;
}

SmallVector<func::FuncOp> GetSortedFunctions(ModuleOp module_op) {
  auto iterator_range = module_op.getOps<func::FuncOp>();
  SmallVector<func::FuncOp> func_ops(iterator_range.begin(),
                                     iterator_range.end());
  absl::c_sort(func_ops, [](func::FuncOp op1, func::FuncOp op2) {
    return op1.getName() < op2.getName();
  });
  return func_ops;
}

// Returns true iff the function's symbol is public.
bool IsPublicFuncOp(func::FuncOp func_op) {
  return SymbolTable::getSymbolVisibility(&*func_op) ==
         SymbolTable::Visibility::Public;
}

}  // namespace

func::FuncOp FindMainFuncOp(ModuleOp module_op) {
  if (const auto main_func_op = module_op.lookupSymbol<func::FuncOp>(
          kImportModelDefaultGraphFuncName);

                               TF::kStablehloEntryFunctionAttrName)
                           .getValue();
  func::FuncOp func_op = symbol_table.lookup<func::FuncOp>(function_name);

  // We should not unwrap if the function is not from
  // ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass.
  if (!func_op->hasAttr(TF::kFromXlaCallModuleAttrName)) {
    return;
  }

  MLIRContext* context = call_op.getContext();

};

void RemoveShardingCustomCallPass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  populateWithGenerated(patterns);

  FrozenRewritePatternSet frozen_patterns(std::move(patterns));
  if (failed(applyPatternsAndFoldGreedily(func_op, frozen_patterns))) {
    func_op.emitWarning() << "Failed to converge "

};

// Folds tf.IdentityOp and tf.IdentityNOp if op device and the argument devices
// from the defining ops match.
void TensorDeviceCopyConversionPass::runOnOperation() {
  func::FuncOp func_op = getOperation();

  auto should_fold_op_func = [&func_op](const Value &arg,
                                        const StringAttr &op_device) {
    // In TFRT TPU, tensor transfer is handled specifically by D2H and

    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) {

// function.
bool ShouldIncludeInMainFunction(func::FuncOp func_op) {
  return !func_op.isPrivate() && IsExported(func_op) &&
         IsEntryFunction(func_op);
}

// Sets a function to be private so it can be referred internally.
void SetFunctionPrivate(func::FuncOp func) {
  func.setVisibility(SymbolTable::Visibility::Private);