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,
       GetInitializerFunctionReturnsNullWhenInitializersEmpty) {

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

// A pass that removes `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)` patterns
// from the initializer function (type = "restore_op").
//

#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::kTfSavedModelIndexPathAttr;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;

// This pass creates a RestoreV2 op in the initializer function with

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

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_saved_model.h"
#include "tensorflow/compiler/mlir/utils/name_utils.h"

namespace mlir {
namespace quant {
namespace {

using ::mlir::tf_saved_model::GetInitializerFunction;
using ::mlir::tf_saved_model::GetSessionInitializerOp;
using ::mlir::tf_saved_model::kTfSavedModelExportedNamesAttr;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;

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()) {

    // `init_func_op` is guaranteed to be not null in a valid module.
    init_func_ops.push_back(init_func_op);
  }

  return init_func_ops;
}

func::FuncOp GetInitializerFunction(ModuleOp module_op,
                                    const StringRef initializer_type) {
  SmallVector<func::FuncOp, 2> init_func_ops =
      GetInitializerFunctions(module_op);