CreateDropWhileShapeInvariantInDeviceClusterPass();

// Creates a pass that moves writes to replicate invariant resource variables
// outside tf_device.replicate op.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateHoistReplicateInvariantResourceWritesPass();

// Transforms functional control flow operations in the TensorFlow dialect to
// MLIR Control Flow Graph (CFG) form.
std::unique_ptr<OperationPass<func::FuncOp>>

    pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
    if (quantization_options.op_set() == OpSet::XLA) {
      pm.addNestedPass<mlir::func::FuncOp>(
          mlir::quant::CreateReplaceCastHacksWithTFXLAOpsPass());
    }
    pm.addNestedPass<mlir::func::FuncOp>(mlir::createCSEPass());
  }
  pm.addNestedPass<mlir::func::FuncOp>(mlir::quant::CreateOptimizePass());
}

    pass_manager.addPass(mlir::TFL::CreatePrepareQuantizeVariablesPass());
    pass_manager.addNestedPass<mlir::func::FuncOp>(
        mlir::TFL::CreateQuantizePass(quant_specs));
    pass_manager.addNestedPass<mlir::func::FuncOp>(
        mlir::TFL::CreatePostQuantizePass(emit_quant_adaptor_ops));
  }
  pass_manager.addNestedPass<mlir::func::FuncOp>(
      mlir::TFL::CreateOptimizeOpOrderPass());

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

// input and output types by unwrapping quantization parameters.
std::unique_ptr<OperationPass<func::FuncOp>> CreateQuantizePass();

// Overloading of CreateQuantizePass which takes QuantizationSpecs.
std::unique_ptr<OperationPass<func::FuncOp>> CreateQuantizePass(
    QuantizationSpecs quant_specs,
    tensorflow::quantization::OpSet target_opset);

// moved to the main function. Returns failure() iff validation fails.
FailureOr<absl::flat_hash_map<std::string, func::FuncOp>> GetInitFuncOps(
    ModuleOp module_op) {
  absl::flat_hash_map<std::string, func::FuncOp> init_func_ops;

  for (func::FuncOp init_func_op : GetInitializerFunctions(module_op)) {
    if (failed(ValidateInitFunc(init_func_op))) {
      return failure();
    }

def QuantizeWeightPass : Pass<"stablehlo-quantize-weight", "mlir::func::FuncOp"> {
  let summary = "Quantizes the weight component of StableHLO graph.";
  let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
  let constructor = "mlir::quant::stablehlo::CreateQuantizeWeightPass()";
}

def UnfuseMhloBatchNormPass : Pass<"stablehlo-unfuse-mhlo-batch-norm", "mlir::func::FuncOp"> {
  let summary = "Unfuses batch normalization into arithmetic ops.";

  // Converts a given FuncOp to a FunctionDef and adds it to the function
  // definition library
  static Status ConvertLibFunction(
      const GraphExportConfig& configs, const Dialect* tf_dialect,
      const SymbolTable& symbol_table, FuncOp function,
      FunctionLibraryDefinition* flib_def,
      llvm::SmallDenseSet<FuncOp>& visited_functions);

TEST_F(LiftAsFunctionCallTest, LiftedFunctionSucceeds) {
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleLifted);
  ASSERT_TRUE(module_op);

  auto composite_dot_general_fn =
      module_op->lookupSymbol<func::FuncOp>("composite_dot_general_fn_1");
  ASSERT_THAT(composite_dot_general_fn, NotNull());

  auto dot_general_op = FindOperationOfType<mlir::stablehlo::DotGeneralOp>(
      composite_dot_general_fn);

  }

  pm.addNestedPass<mlir::func::FuncOp>(mlir::TF::CreateLowerQuantizedPass());
  pm.addNestedPass<mlir::func::FuncOp>(
      mlir::quant::stablehlo::CreateConvertTFQuantTypesPass());

  if (lower_to_xla_hlo) {
    for (auto& target_pass : custom_legalization_passes) {
      pm.addNestedPass<mlir::func::FuncOp>(std::move(target_pass));
    }