"tf_saved_model.session_initializer"() { initializers = [@init] } : () -> ()
  // 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-SAME: exported_names = ["__tf_saved_model_session_initializer"]
  func.func @init(
    %arg0: tensor<!tf_type.string> {tf_saved_model.bound_input = @asset_sym_name},
    %arg1: tensor<!tf_type.resource<tensor<1x64xf32>>> {tf_saved_model.bound_input = @some_constant}
  ) attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer"]}
  {

  "tf_saved_model.session_initializer"() { initializers = [@init] } : () -> ()
  func.func @init() -> (tensor<1xf32> {tf_saved_model.index_path = ["output"]})
    attributes { tf_saved_model.exported_names = ["__tf_saved_model_session_initializer"], tf_saved_model.initializer_type = "init_op" } {
    %0 = "tf.Const"() {value = dense<[1.0]> : tensor<1xf32> } : () -> tensor<1xf32>
    func.return %0 : tensor<1xf32>
  }
}

// -----

module attributes {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"]} {
    %cst = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi64>
    %cst_0 = arith.constant dense<["a", "b", "c", "d"]> : tensor<4x!tf_type.string>

# CHECK:      func @[[restore]](
# CHECK-SAME:   [[variable_path:%.*]]: tensor<!tf_type.string> {tf_saved_model.bound_input = @[[sym_name]]}
# CHECK-SAME: tf_saved_model.exported_names = ["{{__tf_saved_model_session_initializer.*}}"]
# CHECK-SAME: tf_saved_model.initializer_type = "restore_op"
# CHECK: [[v0:%.*]] = "tf.RestoreV2"([[variable_path]]
# CHECK: [[v1:%.*]] = "tf.Identity"([[v0]])
# CHECK: [[handle:%.*]] = "tf.VarHandleOp"

// Test for case with existing session initialize op.
module attributes {tf_saved_model.semantics} {
  // CHECK-LABEL: init_all_tables
  func.func @init_all_tables()
  attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer"]} {
    %cst = arith.constant dense<[1, 2, 3, 4]> : tensor<4xi64>
    %cst_0 = arith.constant dense<["a", "b", "c", "d"]> : tensor<4x!tf_type.string>

  // tf_saved_model.
  init_func_op->setAttr(
      kTfSavedModelExportedNamesAttr,
      builder.getStrArrayAttr({absl::StrCat(
          "__tf_saved_model_session_initializer_", target_node_name)}));
  init_func_op->setAttr(kTfSavedModelInitializerTypeAttr,
                        builder.getStringAttr(initializer_type));