Merge pull request #46716 from yongtang:46676-PopulationCount-uint32-…

…uint64 PiperOrigin-RevId: 359060705 Change-Id: I6cfc38e0a16cb9362b4ea4bcc83fff82fe86c0dd
tensorflow · Feb 23, 2021 · 5b49cfe · 5b49cfe
2 parents fb21da8 + 0118cff
commit 5b49cfe
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/tensorflow/core/kernels/population_count_op.cc b/tensorflow/core/kernels/population_count_op.cc
@@ -65,7 +65,9 @@ TF_CALL_int8(REGISTER_POPULATION_COUNT);
 TF_CALL_uint16(REGISTER_POPULATION_COUNT);
 TF_CALL_int16(REGISTER_POPULATION_COUNT);
 TF_CALL_int32(REGISTER_POPULATION_COUNT);
+TF_CALL_uint32(REGISTER_POPULATION_COUNT);
 TF_CALL_int64(REGISTER_POPULATION_COUNT);
+TF_CALL_uint64(REGISTER_POPULATION_COUNT);
 
 #undef REGISTER_POPULATION_COUNT
 
@@ -87,7 +89,9 @@ POPCNT(uint8, 8);
 POPCNT(int16, 16);
 POPCNT(uint16, 16);
 POPCNT(int32, 32);
+POPCNT(uint32, 32);
 POPCNT(int64, 64);
+POPCNT(uint64, 64);
 
 #undef POPCNT
 

diff --git a/tensorflow/python/ops/bitwise_ops_test.py b/tensorflow/python/ops/bitwise_ops_test.py
@@ -52,9 +52,10 @@ def testBinaryOps(self):
         self.assertAllEqual(xor_result, [5, 5, 4, 5])
 
   def testPopulationCountOp(self):
-    dtype_list = [dtypes.int8, dtypes.int16,
-                  dtypes.int32, dtypes.int64,
-                  dtypes.uint8, dtypes.uint16]
+    dtype_list = [
+        dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64, dtypes.uint8,
+        dtypes.uint16, dtypes.uint32, dtypes.uint64
+    ]
     raw_inputs = [0, 1, -1, 3, -3, 5, -5, 14, -14,
                   127, 128, 255, 256, 65535, 65536,
                   2**31 - 1, 2**31, 2**32 - 1, 2**32, -2**32 + 1, -2**32,