}
           })pb",
      &graph_def));

  // Restore op's name should start with "restore_op" and the save op's name
  // should start with "tf_quant__save_op".
  const std::vector<std::string> control_ret_node_names = {
      "restore_op_0", "tf_quant__save_op_0"};

  TF_ASSERT_OK_AND_ASSIGN(const std::optional<SaverDef> saver_def,

// CHECK: func.func @main(%[[ARG:.*]]: tensor<!tf_type.string> {tf_saved_model.index_path = ["__tf_file_prefix"]})
// CHECK-SAME: {{{.*tf.entry_function = {inputs = "restore_op_0:0", outputs = ""}.*}}}
// CHECK-NEXT: tf_executor.graph

// Checks that the ops from @init_func_restore_op are cloned.

  TF_RETURN_IF_ERROR(restore_op->AddInput(prefix_handle.get()));
  TF_RETURN_IF_ERROR(restore_op->AddInput(names_handle.get()));
  TF_RETURN_IF_ERROR(restore_op->AddInput(shapes_and_slices_handle.get()));

  AbstractTensorHandle* restored_handle = nullptr;
  int num_retvals = 1;
  TF_RETURN_IF_ERROR(restore_op->Execute(
      absl::MakeSpan(&restored_handle, num_retvals), &num_retvals));

A. Unique TensorFlower <******@****.***> 1661895558 -0700

        "//tensorflow/c/experimental/saved_model/core:__subpackages__",
    ],
    licenses = ["notice"],
)

cc_library(
    name = "restore_ops",
    srcs = [
        "restore_ops.cc",
    ],
    hdrs = [
        "restore_ops.h",
    ],
    deps = [
        "//tensorflow/c:tensor_interface",
        "//tensorflow/c/eager:abstract_tensor_handle",

  // All three ops should have been removed.
  // CHECK: @init_func_restore_op
  // CHECK-SAME: tf_saved_model.initializer_type = "restore_op"
  // CHECK-NEXT: return
}

// -----

// The `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)` pattern is not removed
// from the initializer function that is not "restore_op" type.

// CHECK-LABEL: module
module attributes {tf_saved_model.semantics} {

  func.func @init_func_restore_op() -> () attributes {
      tf_saved_model.exported_names = ["tf_saved_model.session_initializer_restore_op"],
      tf_saved_model.initializer_type = "restore_op"} {
    %cst_0 = "tf.Const"() {value = dense<1.000000e+00> : tensor<2xf32>} : () -> tensor<2xf32>
    %var_0 = "tf.VarHandleOp"() {shared_name = "var_0"} : () -> tensor<!tf_type.resource<tensor<2xf32>>>

  // expected-error@+1 {{expects exactly one block in the MLIR function}}
  func.func @init() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer"], tf_saved_model.initializer_type = "restore_op"} {
    cf.br ^bb1
    ^bb1:
    func.return
  }
}

// -----

module attributes {tf_saved_model.semantics} {
  // Test case: No variables
  // CHECK: func @init()
  // CHECK: tf.Const

  auto restore_op = builder.create<TF::RestoreV2Op>(
      session_init_func.getLoc(),
      /*tensors=*/tensor_types,
      /*prefix=*/filename_arg, tensor_names_const, shape_and_slices_const);

  for (auto [idx, restore_result] : llvm::enumerate(restore_op.getResults())) {
    builder.create<TF::AssignVariableOp>(
        restore_op.getLoc(), target_var_handle_ops[idx], restore_result);
  }

  }

  func.func @Init() attributes {tf_saved_model.exported_names = ["Init"], tf_saved_model.initializer_type = "restore_op"} {
    func.return
  }

  // CHECK-LABEL: func @Init
  // CHECK-SAME: tf_saved_model.exported_names = ["Init"]
  // CHECK-SAME: tf_saved_model.initializer_type = "restore_op"
  // CHECK: %[[VAR:.*]] = "tf.VarHandleOp"()
  // CHECK-SAME: "var1"