bool IsEntryFunction(func::FuncOp func);

// Get all the entry functions in an MLIR module.
llvm::SmallVector<func::FuncOp> GetEntryFunctions(ModuleOp module);

// Get all the functions referenced in a symber user op and save them in
// `callees`.
LogicalResult GetCallees(SymbolUserOpInterface op, SymbolTable &symtab,
                         llvm::SmallVector<func::FuncOp> &callees);

  // 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());

  pm.addNestedPass<mlir::func::FuncOp>(
      mlir::mhlo::createLegalizeEinsumToDotGeneralPass());
  pm.addNestedPass<mlir::func::FuncOp>(
      mlir::mhlo::createLegalizeDotToDotGeneralPass());
  // Unfuse mhlo BatchNorm to primitive ops.
  pm.addNestedPass<mlir::func::FuncOp>(mlir::odml::createUnfuseBatchNormPass());
  // Fuse Conv + Mul to Conv.
  pm.addNestedPass<mlir::func::FuncOp>(mlir::odml::createFuseConvolutionPass());

  pm.addNestedPass<mlir::func::FuncOp>(CreateConvertTFQuantOpsToMHLOPass());
  pm.addNestedPass<mlir::func::FuncOp>(mhlo::createChloLegalizeToHloPass());
  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
  pm.addNestedPass<mlir::func::FuncOp>(
      mhlo::createMhloQuantLegalizeToIntPass());
  pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());

  auto func_op = module_op->lookupSymbol<func::FuncOp>("unknown_rank_arg");
  ASSERT_THAT(func_op, NotNull());
  ASSERT_THAT(func_op.getNumArguments(), Eq(1));

  EXPECT_FALSE(HasRankOf(func_op.getArgument(0), /*rank=*/1));
}

TEST_F(AttrsAndConstraintsTest, TryCastSucceeds) {
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleStatic);
  ASSERT_TRUE(module_op);

  func::FuncOp main_fn = FindMainFuncOp(*module_op);

// 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);

  // It's simply clone the FuncOp and hardware specific transformations.
  func::FuncOp GetAlternativeViewForSpec(
      func::FuncOp func,
      const InferenceDeviceType& current_device_inference_type,
      const InferenceDeviceType& target_device_inference_type, ModuleOp module,
      OpBuilder* builder);

  // Apply any device-specific optimizations.
  void Optimize(func::FuncOp func, const std::string& hardware);

// A utility class that enables the conversion of the LSTMCellSimple composite
// op into a fused TFL LSTM op. The fused op is contained within a FuncOp
// that also contains other supporting ops needed to construct the operands for
// the fused op. The caller provides the containing FuncOp as input with
// arguments specifying the input, weight, projection and bias.
// The weight, projection, bias and layer norm scale all need to be

  pm.addNestedPass<FuncOp>(
      mlir::CreateFunctionalToExecutorDialectConversionPass());

  // Do a single pass to split the graph's single island op into an island per
  // op as expected by the following passes.
  pm.addNestedPass<FuncOp>(mlir::TF::CreateSplitIntoIslandPerOpPass());

  pm.addNestedPass<FuncOp>(mlir::TFDevice::CreateReplicateToIslandPass(

          if (func && reachable_funcs.insert(func).second)
            funcs_to_visit.push_back(func);
        });
      };

  SmallVector<func::FuncOp, 4> funcs_to_visit;
  collect_reachable_funcs(region, funcs_to_visit);

  while (!funcs_to_visit.empty()) {
    SmallVector<func::FuncOp, 4> new_funcs_to_visit;
    for (func::FuncOp func_to_visit : funcs_to_visit) {