#include "tsl/platform/logging.h"
#include "tsl/platform/status.h"

namespace tensorflow {
namespace quantization {
namespace {

using ::mlir::func::FuncOp;
using ::mlir::tf_saved_model::GetInitializerFunction;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;

// Adds the tensor that initializes the variable through the provided
// `assign_var_op` to the `bundle_writer` for saving to checkpoint. Returns the

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"

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

// Returns the initializer function whose `tf_saved_model.initializer_type`
// attribute matches `initializer_type`. Returns a null op if it doesn't exist.
func::FuncOp GetInitializerFunction(ModuleOp module_op,
                                    StringRef initializer_type);

// Checks if the module restores variables from a Checkpoint.
bool IsRestoreGraph(ModuleOp module);

func::FuncOp GetOrCreateSessionInitFunc(ModuleOp module) {
  SessionInitializerOp session_init_op = GetSessionInitializerOp(module);
  if (!session_init_op) return CreateSessionInitFunc(module);

  auto init_func_op = GetInitializerFunction(
      module, /*initializer_type=*/kTfSavedModelInitializerRestoreType);
  if (init_func_op) {
    return init_func_op;
  } else if (!session_init_op.getInitializers().empty()) {