namespace mlir {
namespace TFR {

class RewriteQuantizedIOPass
    : public PassWrapper<RewriteQuantizedIOPass, OperationPass<ModuleOp>> {
 public:
  StringRef getArgument() const final { return "tfr-rewrite-quantized-io"; }

  StringRef getDescription() const final {
    return "Replaces operands and results that has quantized type with their "
           "storage types.";
  }
  void runOnOperation() override;
};

==============================================================================*/

#include "tensorflow/compiler/mlir/tfr/utils/utils.h"

#include <string>

#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/StringSet.h"
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"

namespace mlir {
namespace TFR {
namespace {

==============================================================================*/

#include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h"

#include <memory>
#include <string>

#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h"
#include "tensorflow/core/util/debug_data_dumper.h"

namespace tensorflow {

}  // namespace

class DropSavedModelSemanticsPass
    : public PassWrapper<DropSavedModelSemanticsPass, OperationPass<ModuleOp>> {
 public:
  StringRef getArgument() const final { return "drop-savedmodel-semantics"; }
  StringRef getDescription() const final {
    return "Drops all tf_saved_model attributes";
  }

  // NOTE: The implementation is mostly copypasted from

#ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSLATE_EXPORT_TF_DIALECT_OP_H_
#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSLATE_EXPORT_TF_DIALECT_OP_H_

#include "llvm/ADT/StringRef.h"
#include "mlir/IR/Operation.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/utils/export_utils.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"

#include <utility>

#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project

namespace mlir {
namespace quant {
namespace {

using ::mlir::tf_saved_model::GetInitializerFunction;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;

constexpr StringRef kTfQuantSaveV2OpName = "tf_quant__save_save_v2";
constexpr StringRef kTfQuantSaveReturnOpName = "tf_quant__save_return";

// A pass that creates a new function that wraps the newly created SaveV2 op.

// dialect back into a TensorFlow valid op name.
absl::StatusOr<llvm::StringRef> GetTensorFlowOpName(llvm::StringRef);

// Converts an MLIR operation to TensorFlow NodeDef with given node name. This
// name should be unique to the graph it is being inserted into.
absl::StatusOr<std::unique_ptr<NodeDef>> GetOperationNodeDef(
    mlir::Operation* inst, llvm::StringRef name);

// Converts MLIR attributes with values to their tensorflow equivalent.

ABSL_DEPRECATED(
    "Use v2/tf_dialect_to_executor.h::ExportFromTensorflowDialectToExecutor "
    "instead.")
tensorflow::Status ExportFromTensorflowDialectToExecutor(
    mlir::ModuleOp module, llvm::StringRef module_name = llvm::StringRef());

}  // namespace v1
}  // namespace tf2xla
}  // namespace tensorflow

namespace mlir {
namespace odml {

class RenameEntrypointToMainPass
    : public PassWrapper<RenameEntrypointToMainPass, OperationPass<ModuleOp>> {
 public:
  StringRef getArgument() const final { return "rename-entrypoint-to-main"; }
  StringRef getDescription() const final {
    return "Renames the entrypoint in SavedModel to `main`";
  }

  void runOnOperation() override {