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

limitations under the License.
==============================================================================*/

#include "tensorflow/c/experimental/saved_model/core/ops/restore_ops.h"

#include "tensorflow/c/eager/abstract_tensor_handle.h"
#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/test_utils.h"

#include "tensorflow/c/eager/immediate_execution_tensor_handle.h"
#include "tensorflow/c/experimental/saved_model/core/concrete_function.h"
#include "tensorflow/c/experimental/saved_model/core/ops/restore_ops.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/constant.h"
#include "tensorflow/c/experimental/saved_model/core/revived_types/flat_tensor_function.h"

// as saved by STACK_AND_SAVE_HOST_TO_GO_ABI.
//
// R0 is not guaranteed to contain $0 after this macro.
#define UNSTACK_AND_RESTORE_GO_TO_HOST_ABI(extra)                     \
	RESTORE_GPR(extra+FIXED_FRAME)                                \
	RESTORE_FPR(extra+FIXED_FRAME+SAVE_GPR_SIZE)                  \
	RESTORE_VR(extra+FIXED_FRAME+SAVE_GPR_SIZE+SAVE_FPR_SIZE, R0) \
	ADD 	$(extra+FIXED_FRAME+SAVE_ALL_REG_SIZE), R1            \

             }
           })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,

  // 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} {

// CHECK-LOC: tf_executor.island({{.*}}) wraps "tf.NoOp"()
// CHECK-LOC-SAME: loc("init_op_NoOp")
}

// -----

// Test the case where there are 2 initializer functions ("init_op" and
// "restore_op"). The init func of type "init_op" is merged first.

// CHECK-LABEL: module attributes
module attributes {tf.versions = {bad_consumers = [], min_consumer = 12 : i32, producer = 1228 : i32}, 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>>>