entrypoints[op.getSymName()] = op;
    });

    if (auto session_initializer =
            tf_saved_model::GetSessionInitializerOp(module)) {
      // clang-format off
      // Skip session initializer functions which are present in saved models.
      // For example:
      //   "tf_saved_model.session_initializer"() {initializers = [@init_all_tables]} : () -> () // NOLINT

// RUN:   -quant-insert-restore-op | FileCheck %s --check-prefix CHECK-LOC

// RestoreV2 op created for a single VarHandleOp.

module attributes {tf_saved_model.semantics} {
  "tf_saved_model.session_initializer"() {initializers = [@init_func_restore_op]} : () -> ()

  func.func @init_func_restore_op() -> () attributes {
      tf_saved_model.initializer_type = "restore_op",

};

TEST_F(SaveVariablesToCheckpointTest, VariableSavedToCheckpoint) {
  constexpr absl::string_view kModuleCode = R"mlir(
    module attributes {tf_saved_model.semantics} {
      "tf_saved_model.session_initializer"() {initializers = [@init_func_restore_op]} : () -> ()

      func.func @init_func_restore_op() -> () attributes {tf_saved_model.exported_names = ["restore"], tf_saved_model.initializer_type = "restore_op"} {

// RUN: tf-opt %s | tf-opt | FileCheck %s

module attributes {tf_saved_model.semantics} {

  // CHECK: tf_saved_model.session_initializer
  "tf_saved_model.session_initializer"() {
    initializers = [@init]
  } : () -> ()

  // CHECK: tf_saved_model.asset
  "tf_saved_model.asset"() {
    filename = "asset_filename",
    sym_name = "asset_sym_name"
  } : () -> ()

  // Representation for constants: (immutable) global tensor.

  "tf_saved_model.session_initializer"() {initializers = [@NoOp]} : () -> ()
  "tf_saved_model.asset"() {filename = "assets/mydata.txt", sym_name = "__tf_saved_model_asset0_mydata.txt"} : () -> ()
// Session initializer ops and asset ops untouched.
// CHECK: "tf_saved_model.session_initializer"() <{initializers = [@NoOp]}> : () -> ()

  session_init_op.setInitializersAttr(
      ArrayAttr::get(session_init_op.getContext(), initializers_attrs));
}

// Returns the session_initializer op in the module if exists. Otherwise,
// creates a new session_initializer op and returns it.
SessionInitializerOp GetOrCreateSessionInitializerOp(ModuleOp module_op) {
  SessionInitializerOp session_init_op = GetSessionInitializerOp(module_op);

// CHECK: tfrt.cost_threshold = 1024
// CHECK-SAME: tfrt.merge_inter_dependent_streams = true
module attributes {tf_saved_model.semantics} {

// CHECK-NOT: "tf_saved_model.session_initializer"
"tf_saved_model.session_initializer"() { initializers = [@func_init] } : () -> ()

// CHECK-LABEL: _tfrt_resource_init
// CHECK: tf.VarHandleOp
// CHECK: tf.ReadVariableOp
// CHECK: tfrt_fallback_async.set_resource

// RUN: tf-opt %s -split-input-file -tf-saved-model-asset-sinking='saved-model-dir=foo/bar' | FileCheck %s

// CHECK-LABEL: module @asset
module @asset attributes {tf_saved_model.semantics} {
  "tf_saved_model.session_initializer"() {initializers = [@init]} : () -> ()

  // CHECK-NOT: "tf_saved_model.asset"
  "tf_saved_model.asset"() {filename = "assets/test0.txt", sym_name = "asset0"} : () -> ()

}

// -----

// Test with an existing session init function.
module attributes {tf_saved_model.semantics, tf_saved_model.under_construction} {

  "tf_saved_model.session_initializer"() {initializers = [@Init]} : () -> ()
  func.func @serving_default(%arg0: tensor<!tf_type.resource<tensor<100x50xf32>>> {tf.resource_name = "dense/kernel"}) -> (tensor<100x50xf32> {tf_saved_model.index_path = ["dense_2"]})

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1506 : i32}, tf_saved_model.semantics} {
  "tf_saved_model.session_initializer"() {initializers = [@init_all_tables]} : () -> ()
  func.func @init_all_tables() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer_init_all_tables"], tf_saved_model.initializer_type = "init_op"} {