Add int64-int64 key-value in combination on GPU.

Author: Lifann

tensorflow_recommenders_addons/dynamic_embedding/core/kernels/cuckoo_hashtable_op_gpu.cu.cc

@@ -662,6 +662,7 @@ REGISTER_KERNEL_BUILDER(
 
 REGISTER_KERNEL(int64, float);
 REGISTER_KERNEL(int64, Eigen::half);
+REGISTER_KERNEL(int64, int64);
 REGISTER_KERNEL(int64, int32);
 REGISTER_KERNEL(int64, int8);
 

tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_gpu.h

@@ -176,6 +176,7 @@ void CreateTableImpl(TableWrapperBase<K, V>** pptable, size_t max_size,
 
 DECLARE_CREATE_TABLE(int64, float);
 DECLARE_CREATE_TABLE(int64, Eigen::half);
+DECLARE_CREATE_TABLE(int64, int64);
 DECLARE_CREATE_TABLE(int64, int32);
 DECLARE_CREATE_TABLE(int64, int8);
 

tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_gpu_impl_int64_0.cu.cc

@@ -0,0 +1,25 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_gpu.h"
+namespace tensorflow {
+namespace recommenders_addons {
+namespace lookup {
+namespace gpu {
+DEFINE_CREATE_TABLE(0, int64, int64, 0, 0);
+}  // namespace gpu
+}  // namespace lookup
+}  // namespace recommenders_addons
+}  // namespace tensorflow

tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_gpu_impl_int64_1.cu.cc

@@ -0,0 +1,25 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_gpu.h"
+namespace tensorflow {
+namespace recommenders_addons {
+namespace lookup {
+namespace gpu {
+DEFINE_CREATE_TABLE(1, int64, int64, 0, 5);
+}  // namespace gpu
+}  // namespace lookup
+}  // namespace recommenders_addons
+}  // namespace tensorflow

tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_gpu_impl_int64_2.cu.cc

@@ -0,0 +1,25 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_gpu.h"
+namespace tensorflow {
+namespace recommenders_addons {
+namespace lookup {
+namespace gpu {
+DEFINE_CREATE_TABLE(2, int64, int64, 1, 0);
+}  // namespace gpu
+}  // namespace lookup
+}  // namespace recommenders_addons
+}  // namespace tensorflow

tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_gpu_impl_int64_3.cu.cc

@@ -0,0 +1,25 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow_recommenders_addons/dynamic_embedding/core/kernels/lookup_impl/lookup_table_op_gpu.h"
+namespace tensorflow {
+namespace recommenders_addons {
+namespace lookup {
+namespace gpu {
+DEFINE_CREATE_TABLE(3, int64, int64, 1, 5);
+}  // namespace gpu
+}  // namespace lookup
+}  // namespace recommenders_addons
+}  // namespace tensorflow

tensorflow_recommenders_addons/dynamic_embedding/python/kernel_tests/dynamic_embedding_variable_test.py

@@ -370,7 +370,7 @@ def test_variable(self):
       dim_list = [1, 2, 4, 8, 10, 16, 32, 64, 100, 200]
       kv_list = [[dtypes.int64, dtypes.float32], [dtypes.int64, dtypes.int32],
                  [dtypes.int64, dtypes.half], [dtypes.int64, dtypes.int8],
-                 [dtypes.int32, dtypes.float32]]
+                 [dtypes.int32, dtypes.float32], [dtypes.int64, dtypes.int64]]
     else:
       dim_list = [1, 8, 16, 128]
       kv_list = [[dtypes.int32, dtypes.double], [dtypes.int32, dtypes.float32],