: OpRewritePattern<TF::XlaCallModuleOp>(context), module_op_(module_op) {}
  using OpRewritePattern<TF::XlaCallModuleOp>::OpRewritePattern;

 private:
  ModuleOp module_op_;
  LogicalResult matchAndRewrite(TF::XlaCallModuleOp op,
                                PatternRewriter &rewriter) const override {
    OwningOpRef<ModuleOp> stablehlo_module_op =

 public:
  explicit ConvertCustomCallWithApproxTopK(MLIRContext* context,
                                           mlir::ModuleOp* module_op)
      : OpConversionPattern<mhlo::CustomCallOp>(context, /*benefit=*/0),
        module_op_(module_op) {}

  mlir::LogicalResult matchAndRewrite(
      mhlo::CustomCallOp op, OpAdaptor adaptor,
      ConversionPatternRewriter& rewriter) const final {

  OwningOpRef<ModuleOp> module_op =
      ParseModuleOpString(kConstantOpWithRankFour);
  ASSERT_TRUE(module_op);

  ASSERT_FALSE(module_op->getBodyRegion().empty());
  ASSERT_FALSE(module_op->getBodyRegion().front().empty());

  auto constant_op = dyn_cast_or_null<mlir::stablehlo::ConstantOp>(
      module_op->getBodyRegion().front().front());
  ASSERT_THAT(constant_op, NotNull());

};

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) {
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
    module {

  // MLIR @main function corresponds to the TF function "main_original".
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(R"mlir(
    func.func private @main(%arg: tensor<1x2xf32>) -> (tensor<1x2xf32>) attributes {tf._original_func_name = "main_original"} {
      return %arg : tensor<1x2xf32>
    }
  )mlir");
  ASSERT_TRUE(module_op);

  absl::flat_hash_map<FunctionName, FunctionAlias> function_aliases;

  constexpr StringRef kModuleOpStr =
      R"mlir(module attributes {tf_saved_model.semantics} {})mlir";

  Block block;
  ModuleOp module_op = ParseModuleOp(kModuleOpStr, block, ctx_);

  func::FuncOp init_func_op = GetInitializerFunction(
      module_op, /*initializer_type=*/kTfSavedModelInitializerInitType);

  EXPECT_THAT(init_func_op, IsNull());
}

TEST_F(TfSavedModelTest,

    module {
      func.func @main() -> () {
        return
      }
    }
  )mlir";

  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleWithMainFunc);
  EXPECT_THAT(*module_op, NotNull());

  func::FuncOp main_func_op = FindMainFuncOp(*module_op);
  EXPECT_THAT(main_func_op, NotNull());
}

TEST_F(FindMainFuncOpTest, ReturnsNullWhenMainFuncOpIsPrivate) {

      return %0 : tensor<1x3xf32>
    }
  }
)mlir";

TEST_F(IsOpQuantizableStableHloTest, ConstantOpQuantizable) {
  OwningOpRef<ModuleOp> module_op = ParseModuleOpString(kModuleConstantAdd);
  ASSERT_TRUE(module_op);

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

  auto constant_op =

absl::StatusOr<ExportedModel> ExportDebuggingModel(
    mlir::ModuleOp module_op, mlir::MLIRContext *context,
    const QuantizationOptions &quantization_options,
    const absl::flat_hash_map<std::string, std::string> &function_aliases) {
  // Clone ModuleOp and function aliases so changes in this pipeline won't
  // be reflected in the original values.
  mlir::OwningOpRef<mlir::ModuleOp> cloned_module_ref(module_op.clone());

// passed to the target function.
LogicalResult LiftHashTableOpsToArguments(ModuleOp module_op,
                                          func::FuncOp target_func) {
  if (!llvm::hasSingleElement(target_func)) return success();
  if (!UsedBySupportedOps(module_op, target_func)) return success();
  if (IsMainOrInitializerFunction(module_op, target_func)) return success();

  llvm::StringMap<int> shared_name_to_arg_idx;