func.return %0 : tensor<i32>
}
)mlir";
  mlir::MLIRContext context;
  context.loadDialect<mlir::func::FuncDialect, mlir::TF::TensorFlowDialect>();
  mlir::OwningOpRef<mlir::ModuleOp> module =
      mlir::parseSourceString<mlir::ModuleOp>(code, &context);
  ASSERT_TRUE(module);
  EXPECT_TRUE(IsSupportedByReplicatedBridge(*module));
}

TEST(HasTPUPartitionedCallOpInModule, HasTPUPartitionedCallModule) {

//
// The constants whose sizes are smaller than `size_threshold_in_bytes_` will
// not be converted to variables.
class UnfreezeConstantsPass
    : public PassWrapper<UnfreezeConstantsPass, OperationPass<ModuleOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(UnfreezeConstantsPass)

  explicit UnfreezeConstantsPass()
      : UnfreezeConstantsPass(kDefaultConstantSizeThresholdInBytes) {}

#include "tensorflow/compiler/mlir/quantization/tensorflow/passes/tf_quant_ops.h"

namespace stablehlo::quantization {

void DisableDebugging(mlir::ModuleOp module_op) {
  module_op.walk(
      [](mlir::TF::DumpTensorOp dump_op) { dump_op.setEnabled(false); });
}

void ChangeToQuantizedFilename(mlir::ModuleOp module_op) {
  module_op.walk([](mlir::TF::DumpTensorOp dump_op) {
    dump_op.setFileName("quantized_tensor_data.pb");
  });
}

    mlir::ModuleOp module, const GraphExportConfig& configs);

// Converts an MLIR module to TensorFlow graph and FunctionLibraryDefinition.
// The "main" function of the module is stored in the graph and the rest of
// functions are stored in the library. Control ret nodes are stored separately
// in `control_ret_nodes`.
Status ConvertMlirToGraph(mlir::ModuleOp module,

                                       /*RequiresNullTerminator=*/false);

  llvm::SourceMgr source_mgr;
  source_mgr.AddNewSourceBuffer(std::move(mem_buffer), llvm::SMLoc());
  OwningOpRef<ModuleOp> module_ref =
      parseSourceFile<ModuleOp>(source_mgr, module.getContext());
  // Inline and optimize loaded functions.
  MLIRContext* context = &getContext();
  PassManager pm(context);
  pm.addPass(createInlinerPass());

  }

  llvm::SourceMgr source_mgr;
  source_mgr.AddNewSourceBuffer(std::move(file), llvm::SMLoc());
  return OwningOpRef<mlir::ModuleOp>(
      parseSourceFile<mlir::ModuleOp>(source_mgr, context));
}

TEST(ErrorCollectorTest, TessSuccessPass) {
  std::string input_file = tensorflow::GetDataDependencyFilepath(
      "tensorflow/compiler/mlir/lite/metrics/testdata/strided_slice.mlir");
  MLIRContext context;

    mlir_module_ =
        mlir::parseSourceFile<mlir::ModuleOp>(mlir_module_path, &context_);
    if (!mlir_module_) {
      return absl::Status(
          absl::StatusCode::kNotFound,
          absl::StrCat("Could not find MLIR module at ", mlir_module_path));
    }
    return absl::OkStatus();
  }

  DialectRegistry registry_;
  MLIRContext context_;
  OwningOpRef<mlir::ModuleOp> mlir_module_;
};

        absl::StrCat("Failed to open input file: ", error_message));
  }

  if (input_mlir) {
    source_mgr->AddNewSourceBuffer(std::move(file), llvm::SMLoc());
    return OwningOpRef<ModuleOp>(
        mlir::parseSourceFile<mlir::ModuleOp>(*source_mgr, context));
  }

  // Register extra TF ops passed as OpDef.
  auto extra_opdefs_status = RegisterExtraTfOpDefs(extra_tf_opdefs);

// function of type "restore_op" to identify the VarHandleOps that should be
// saved using the SaveV2 op.
class InsertSaveOpPass
    : public PassWrapper<InsertSaveOpPass, OperationPass<ModuleOp>> {
 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(InsertSaveOpPass)

  explicit InsertSaveOpPass() = default;

  // The argument used to refer to the pass in the textual format (e.g. on the

    // Whether the input file is an MLIR not tflite file.
    bool input_mlir = false;
  };

  explicit TfLiteImporter(const Options& options) : options_(options) {}

  absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> Import() override;

 private:
  Options options_;
  mlir::MLIRContext context_;
  llvm::SourceMgr source_mgr_;
  std::unique_ptr<mlir::SourceMgrDiagnosticHandler> source_mgr_handler_;
};