return %1 : tensor<*xi64>
  }

// Check that the caller is updated.
// CHECK: func.func @main
// CHECK: %[[HASHTABLE_1:.*]] = "tf.HashTableV2"()
// CHECK: %[[HASHTABLE_2:.*]] = "tf.HashTableV2"()
// CHECK: %[[OUT_2:.*]] = "tf.PartitionedCall"(%arg0, %[[HASHTABLE_1]], %[[HASHTABLE_2]])

# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

# RUN: %p/hash_table_v1 | FileCheck %s

# pylint: disable=missing-docstring,line-too-long
import tensorflow.compat.v1 as tf
from tensorflow.compiler.mlir.tensorflow.tests.tf_saved_model import common_v1

// HashTable is initialized if it appears in the initializer since it can't
// check the actual value.
bool IsResourceInitialized(ModuleOp module_op, Operation* hash_table) {
  StringRef shared_name = GetSharedName(hash_table);
  if (shared_name.empty()) return false;

  for (func::FuncOp init_func_op :
       tf_saved_model::GetInitializerFunctions(module_op)) {

func.func @init_all_tables() {
  %cst = arith.constant dense<"tokens.txt"> : tensor<!tf_type.string>
  %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "hash_table_/tmp/vocab.txt_-2_-1", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>

func.func @init_all_tables() {
  %cst = arith.constant dense<"vocab_file_does_not_exist.txt"> : tensor<!tf_type.string>
  %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "hash_table_/tmp/vocab.txt_-2_-1", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>

    // CHECK: [[HASHTABLE:%.+]] = "tf.HashTableV2"()
    // CHECK: "tf.InitializeTableFromTextFileV2"([[HASHTABLE]], [[CST]])
    // CHECK: [[HASHTABLE_1:%.+]] = "tf.HashTableV2"()
    // CHECK: "tf.InitializeTableFromTextFileV2"([[HASHTABLE_1]], [[CST_1]])
    func.return
  }
}

// -----

// Test running the pass on a module that does not have
// tf_saved_model.semantics.

  // CHECK-NEXT:  [[HASH_TABLE:%.*]] = "tfl.hashtable"() <{key_dtype = !tf_type.string, table_id = -1323619995 : i32, value_dtype = i64}> : () -> tensor<1x!tf_type.resource>
  // CHECK-NEXT:   "tfl.hashtable_import"([[HASH_TABLE]], [[CST]], [[CST_1]]) : (tensor<1x!tf_type.resource>, tensor<3x!tf_type.string>, tensor<3xi64>) -> ()

func.func @init_all_tables() {
  %cst = arith.constant dense<"%FILE_PLACEHOLDER"> : tensor<!tf_type.string>
  %0 = "tf.HashTableV2"() {container = "", device = "", key_dtype = !tf_type.string, shared_name = "hash_table_/tmp/vocab.txt_-2_-1", use_node_name_sharing = false, value_dtype = i64} : () -> tensor<!tf_type.resource>

py_strict_binary(
    name = "keras",
    srcs = ["keras.py"],
    deps = [
        ":common",
        "//tensorflow:tensorflow_py",
    ],
)

py_strict_binary(
    name = "hash_table_v1",
    srcs = ["hash_table_v1.py"],
    deps = [
        ":common_v1",
        "//tensorflow:tensorflow_py",
    ],
)

py_strict_binary(
    name = "include_variables_in_init_v1",

    return
  }
  // CHECK: @init_func_restore_op
  // CHECK-SAME: tf_saved_model.initializer_type = "restore_op"
  // CHECK: %[[CST:.*]] = "tf.Const"()
  // CHECK-NEXT: %[[HASH_TABLE:.*]] = "tf.HashTableV2"()
  // CHECK-NEXT: "tf.AssignVariableOp"(%[[HASH_TABLE]], %[[CST]])
}

// -----


// Nothing happens when there are no `tf_saved_model.session_initializer`.

// CHECK-LABEL: module