CVE-2022-21734 (published 2022-02-03T13:15Z)

Tensorflow is an Open Source Machine Learning Framework. The implementation of `MapStage` is vulnerable a `CHECK`-fail if the key tensor is not a scalar. The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, as these are also affected and still in supported range.

Severity: MEDIUM
CVSS v2 base score: 4.0 (AV:N/AC:L/Au:S/C:N/I:N/A:P)
CVSS v3 base score: 6.5, MEDIUM (CVSS:3.1/AV:N/AC:L/PR:L/UI:N/S:U/C:N/I:N/A:H)

Weakness information:
cwe_id: CWE-843
cwe_name: Access of Resource Using Incompatible Type ('Type Confusion')
description: The product allocates or initializes a resource such as a pointer, object, or variable using one type, but it later accesses that resource using a type that is incompatible with the original type.
extended_description: 
            When the product accesses the resource using an incompatible type, this could trigger logical errors because the resource does not have expected properties. In languages without memory safety, such as C and C++, type confusion can lead to out-of-bounds memory access.
            While this weakness is frequently associated with unions when parsing data with many different embedded object types in C, it can be present in any application that can interpret the same variable or memory location in multiple ways.
            This weakness is not unique to C and C++. For example, errors in PHP applications can be triggered by providing array parameters when scalars are expected, or vice versa. Languages such as Perl, which perform automatic conversion of a variable of one type when it is accessed as if it were another type, can also contain these issues.
         
         

Fix information:
repo_url: https://github.com/tensorflow/tensorflow
commit_details: repo_url: https://github.com/tensorflow/tensorflow
author: Isha Arkatkar
author_date: 2021-12-02 18:57:58-08:00
msg: Add a check for Key being scalar tensor for MapStage and OrderedMapStage ops.

According to documentation[1][2], key must be int64 value, but this wasn't enforced and the ops would fail with check failure for non-scalar key value.

[1]https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/ordered-map-stage
[2]https://www.tensorflow.org/api_docs/cc/class/tensorflow/ops/map-stage

PiperOrigin-RevId: 413822112
Change-Id: I9d118faf990e6361900aa32272eff486ad9f0e2e
repository: repo_url: https://github.com/tensorflow/tensorflow
repo_name: tensorflow/tensorflow
description: An Open Source Machine Learning Framework for Everyone
repo_language: C++
owner: tensorflow
file_changes: filename: map_stage_op.cc
change_type: ModificationType.MODIFY
diff: @@ -536,6 +536,11 @@ class MapStageOp : public OpKernel {
     OP_REQUIRES(ctx, key_tensor->NumElements() > 0,
                 errors::InvalidArgument("key must not be empty"));
 
+    OP_REQUIRES(ctx, key_tensor->NumElements() == 1,
+                errors::InvalidArgument(
+                    "key must be an int64 scalar, got tensor with shape: ",
+                    key_tensor->shape()));
+
     // Create copy for insertion into Staging Area
     Tensor key(*key_tensor);
 

programming_language: C++
filename: map_stage_op_test.py
change_type: ModificationType.MODIFY
diff: @@ -12,8 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from tensorflow.python.framework import errors
+import numpy as np
+
+from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
 from tensorflow.python.framework import ops
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
@@ -28,7 +31,7 @@ class MapStageTest(test.TestCase):
 
   @test_util.run_deprecated_v1
   def testSimple(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         pi = array_ops.placeholder(dtypes.int64)
@@ -40,9 +43,9 @@ class MapStageTest(test.TestCase):
         k, y = stager.get(gi)
         y = math_ops.reduce_max(math_ops.matmul(y, y))
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       sess.run(stage, feed_dict={x: -1, pi: 0})
       for i in range(10):
         _, yval = sess.run([stage, y], feed_dict={x: i, pi: i + 1, gi: i})
@@ -50,7 +53,7 @@ class MapStageTest(test.TestCase):
 
   @test_util.run_deprecated_v1
   def testMultiple(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         pi = array_ops.placeholder(dtypes.int64)
@@ -62,9 +65,9 @@ class MapStageTest(test.TestCase):
         k, (z, y) = stager.get(gi)
         y = math_ops.reduce_max(z * math_ops.matmul(y, y))
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       sess.run(stage, feed_dict={x: -1, pi: 0})
       for i in range(10):
         _, yval = sess.run([stage, y], feed_dict={x: i, pi: i + 1, gi: i})
@@ -73,26 +76,25 @@ class MapStageTest(test.TestCase):
 
   @test_util.run_deprecated_v1
   def testDictionary(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         pi = array_ops.placeholder(dtypes.int64)
         gi = array_ops.placeholder(dtypes.int64)
         v = 2. * (array_ops.zeros([128, 128]) + x)
       with ops.device(test.gpu_device_name()):
-        stager = data_flow_ops.MapStagingArea(
-            [dtypes.float32, dtypes.float32],
-            shapes=[[], [128, 128]],
-            names=['x', 'v'])
+        stager = data_flow_ops.MapStagingArea([dtypes.float32, dtypes.float32],
+                                              shapes=[[], [128, 128]],
+                                              names=['x', 'v'])
         stage = stager.put(pi, {'x': x, 'v': v})
         key, ret = stager.get(gi)
         z = ret['x']
         y = ret['v']
         y = math_ops.reduce_max(z * math_ops.matmul(y, y))
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       sess.run(stage, feed_dict={x: -1, pi: 0})
       for i in range(10):
         _, yval = sess.run([stage, y], feed_dict={x: i, pi: i + 1, gi: i})
@@ -102,7 +104,7 @@ class MapStageTest(test.TestCase):
   def testColocation(self):
     gpu_dev = test.gpu_device_name()
 
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         v = 2. * (array_ops.zeros([128, 128]) + x)
@@ -119,58 +121,56 @@ class MapStageTest(test.TestCase):
         self.assertEqual(y.device, '/device:CPU:0')
         self.assertEqual(z[0].device, '/device:CPU:0')
 
-    G.finalize()
+    g.finalize()
 
   @test_util.run_deprecated_v1
   def testPeek(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.int32, name='x')
         pi = array_ops.placeholder(dtypes.int64)
         gi = array_ops.placeholder(dtypes.int64)
         p = array_ops.placeholder(dtypes.int32, name='p')
       with ops.device(test.gpu_device_name()):
-        stager = data_flow_ops.MapStagingArea(
-            [
-                dtypes.int32,
-            ], shapes=[[]])
+        stager = data_flow_ops.MapStagingArea([
+            dtypes.int32,
+        ], shapes=[[]])
         stage = stager.put(pi, [x], [0])
         peek = stager.peek(gi)
         size = stager.size()
 
-    G.finalize()
+    g.finalize()
 
     n = 10
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       for i in range(n):
         sess.run(stage, feed_dict={x: i, pi: i})
 
       for i in range(n):
-        self.assertTrue(sess.run(peek, feed_dict={gi: i})[0] == i)
+        self.assertEqual(sess.run(peek, feed_dict={gi: i})[0], i)
 
-      self.assertTrue(sess.run(size) == 10)
+      self.assertEqual(sess.run(size), 10)
 
   @test_util.run_deprecated_v1
   def testSizeAndClear(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32, name='x')
         pi = array_ops.placeholder(dtypes.int64)
         gi = array_ops.placeholder(dtypes.int64)
         v = 2. * (array_ops.zeros([128, 128]) + x)
       with ops.device(test.gpu_device_name()):
-        stager = data_flow_ops.MapStagingArea(
-            [dtypes.float32, dtypes.float32],
-            shapes=[[], [128, 128]],
-            names=['x', 'v'])
+        stager = data_flow_ops.MapStagingArea([dtypes.float32, dtypes.float32],
+                                              shapes=[[], [128, 128]],
+                                              names=['x', 'v'])
         stage = stager.put(pi, {'x': x, 'v': v})
         size = stager.size()
         clear = stager.clear()
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       sess.run(stage, feed_dict={x: -1, pi: 3})
       self.assertEqual(sess.run(size), 1)
       sess.run(stage, feed_dict={x: -1, pi: 1})
@@ -182,22 +182,23 @@ class MapStageTest(test.TestCase):
   def testCapacity(self):
     capacity = 3
 
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.int32, name='x')
         pi = array_ops.placeholder(dtypes.int64, name='pi')
         gi = array_ops.placeholder(dtypes.int64, name='gi')
       with ops.device(test.gpu_device_name()):
-        stager = data_flow_ops.MapStagingArea(
-            [
-                dtypes.int32,
-            ], capacity=capacity, shapes=[[]])
+        stager = data_flow_ops.MapStagingArea([
+            dtypes.int32,
+        ],
+                                              capacity=capacity,
+                                              shapes=[[]])
 
       stage = stager.put(pi, [x], [0])
       get = stager.get()
       size = stager.size()
 
-    G.finalize()
+    g.finalize()
 
     from six.moves import queue as Queue
     import threading
@@ -205,7 +206,7 @@ class MapStageTest(test.TestCase):
     queue = Queue.Queue()
     n = 8
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       # Stage data in a separate thread which will block
       # when it hits the staging area's capacity and thus
       # not fill the queue with n tokens
@@ -234,13 +235,13 @@ class MapStageTest(test.TestCase):
                                              capacity))
 
       # Should have capacity elements in the staging area
-      self.assertTrue(sess.run(size) == capacity)
+      self.assertEqual(sess.run(size), capacity)
 
       # Clear the staging area completely
       for i in range(n):
         sess.run(get)
 
-      self.assertTrue(sess.run(size) == 0)
+      self.assertEqual(sess.run(size), 0)
 
   @test_util.run_deprecated_v1
   def testMemoryLimit(self):
@@ -248,28 +249,28 @@ class MapStageTest(test.TestCase):
     chunk = 200 * 1024  # 256K
     capacity = memory_limit // chunk
 
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.uint8, name='x')
         pi = array_ops.placeholder(dtypes.int64, name='pi')
         gi = array_ops.placeholder(dtypes.int64, name='gi')
       with ops.device(test.gpu_device_name()):
-        stager = data_flow_ops.MapStagingArea(
-            [dtypes.uint8], memory_limit=memory_limit, shapes=[[]])
+        stager = data_flow_ops.MapStagingArea([dtypes.uint8],
+                                              memory_limit=memory_limit,
+                                              shapes=[[]])
         stage = stager.put(pi, [x], [0])
         get = stager.get()
         size = stager.size()
 
-    G.finalize()
+    g.finalize()
 
     from six.moves import queue as Queue
     import threading
-    import numpy as np
 
     queue = Queue.Queue()
     n = 8
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       # Stage data in a separate thread which will block
       # when it hits the staging area's capacity and thus
       # not fill the queue with n tokens
@@ -299,56 +300,57 @@ class MapStageTest(test.TestCase):
                                              capacity))
 
       # Should have capacity elements in the staging area
-      self.assertTrue(sess.run(size) == capacity)
+      self.assertEqual(sess.run(size), capacity)
 
       # Clear the staging area completely
       for i in range(n):
         sess.run(get)
 
-      self.assertTrue(sess.run(size) == 0)
+      self.assertEqual(sess.run(size), 0)
 
   @test_util.run_deprecated_v1
   def testOrdering(self):
     import six
     import random
 
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.int32, name='x')
         pi = array_ops.placeholder(dtypes.int64, name='pi')
         gi = array_ops.placeholder(dtypes.int64, name='gi')
       with ops.device(test.gpu_device_name()):
-        stager = data_flow_ops.MapStagingArea(
-            [
-                dtypes.int32,
-            ], shapes=[[]], ordered=True)
+        stager = data_flow_ops.MapStagingArea([
+            dtypes.int32,
+        ],
+                                              shapes=[[]],
+                                              ordered=True)
         stage = stager.put(pi, [x], [0])
         get = stager.get()
         size = stager.size()
 
-    G.finalize()
+    g.finalize()
 
     n = 10
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       # Keys n-1..0
       keys = list(reversed(six.moves.range(n)))
 
       for i in keys:
         sess.run(stage, feed_dict={pi: i, x: i})
 
-      self.assertTrue(sess.run(size) == n)
+      self.assertEqual(sess.run(size), n)
 
       # Check that key, values come out in ascending order
       for i, k in enumerate(reversed(keys)):
         get_key, values = sess.run(get)
         self.assertTrue(i == k == get_key == values)
 
-      self.assertTrue(sess.run(size) == 0)
+      self.assertEqual(sess.run(size), 0)
 
   @test_util.run_deprecated_v1
   def testPartialDictInsert(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         f = array_ops.placeholder(dtypes.float32)
@@ -366,41 +368,39 @@ class MapStageTest(test.TestCase):
         size = stager.size()
         isize = stager.incomplete_size()
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       # 0 complete and incomplete entries
-      self.assertTrue(sess.run([size, isize]) == [0, 0])
+      self.assertEqual(sess.run([size, isize]), [0, 0])
       # Stage key 0, x and f tuple entries
       sess.run(stage_xf, feed_dict={pi: 0, x: 1, f: 2})
-      self.assertTrue(sess.run([size, isize]) == [0, 1])
+      self.assertEqual(sess.run([size, isize]), [0, 1])
       # Stage key 1, x and f tuple entries
       sess.run(stage_xf, feed_dict={pi: 1, x: 1, f: 2})
-      self.assertTrue(sess.run([size, isize]) == [0, 2])
+      self.assertEqual(sess.run([size, isize]), [0, 2])
 
       # Now complete key 0 with tuple entry v
       sess.run(stage_v, feed_dict={pi: 0, v: 1})
       # 1 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [1, 1])
+      self.assertEqual(sess.run([size, isize]), [1, 1])
       # We can now obtain tuple associated with key 0
-      self.assertTrue(
-          sess.run([key, ret], feed_dict={
-              gi: 0
-          }) == [0, {
+      self.assertEqual(
+          sess.run([key, ret], feed_dict={gi: 0}),
+          [0, {
               'x': 1,
               'f': 2,
               'v': 1
           }])
 
       # 0 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [0, 1])
+      self.assertEqual(sess.run([size, isize]), [0, 1])
       # Now complete key 1 with tuple entry v
       sess.run(stage_v, feed_dict={pi: 1, v: 3})
       # We can now obtain tuple associated with key 1
-      self.assertTrue(
-          sess.run([key, ret], feed_dict={
-              gi: 1
-          }) == [1, {
+      self.assertEqual(
+          sess.run([key, ret], feed_dict={gi: 1}),
+          [1, {
               'x': 1,
               'f': 2,
               'v': 3
@@ -408,7 +408,7 @@ class MapStageTest(test.TestCase):
 
   @test_util.run_deprecated_v1
   def testPartialIndexInsert(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         f = array_ops.placeholder(dtypes.float32)
@@ -424,35 +424,35 @@ class MapStageTest(test.TestCase):
         size = stager.size()
         isize = stager.incomplete_size()
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       # 0 complete and incomplete entries
-      self.assertTrue(sess.run([size, isize]) == [0, 0])
+      self.assertEqual(sess.run([size, isize]), [0, 0])
       # Stage key 0, x and f tuple entries
       sess.run(stage_xf, feed_dict={pi: 0, x: 1, f: 2})
-      self.assertTrue(sess.run([size, isize]) == [0, 1])
+      self.assertEqual(sess.run([size, isize]), [0, 1])
       # Stage key 1, x and f tuple entries
       sess.run(stage_xf, feed_dict={pi: 1, x: 1, f: 2})
-      self.assertTrue(sess.run([size, isize]) == [0, 2])
+      self.assertEqual(sess.run([size, isize]), [0, 2])
 
       # Now complete key 0 with tuple entry v
       sess.run(stage_v, feed_dict={pi: 0, v: 1})
       # 1 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [1, 1])
+      self.assertEqual(sess.run([size, isize]), [1, 1])
       # We can now obtain tuple associated with key 0
-      self.assertTrue(sess.run([key, ret], feed_dict={gi: 0}) == [0, [1, 1, 2]])
+      self.assertEqual(sess.run([key, ret], feed_dict={gi: 0}), [0, [1, 1, 2]])
 
       # 0 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [0, 1])
+      self.assertEqual(sess.run([size, isize]), [0, 1])
       # Now complete key 1 with tuple entry v
       sess.run(stage_v, feed_dict={pi: 1, v: 3})
       # We can now obtain tuple associated with key 1
-      self.assertTrue(sess.run([key, ret], feed_dict={gi: 1}) == [1, [1, 3, 2]])
+      self.assertEqual(sess.run([key, ret], feed_dict={gi: 1}), [1, [1, 3, 2]])
 
   @test_util.run_deprecated_v1
   def testPartialDictGetsAndPeeks(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         f = array_ops.placeholder(dtypes.float32)
@@ -476,40 +476,38 @@ class MapStageTest(test.TestCase):
         size = stager.size()
         isize = stager.incomplete_size()
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       # 0 complete and incomplete entries
-      self.assertTrue(sess.run([size, isize]) == [0, 0])
+      self.assertEqual(sess.run([size, isize]), [0, 0])
       # Stage key 0, x and f tuple entries
       sess.run(stage_xf, feed_dict={pi: 0, x: 1, f: 2})
-      self.assertTrue(sess.run([size, isize]) == [0, 1])
+      self.assertEqual(sess.run([size, isize]), [0, 1])
       # Stage key 1, x and f tuple entries
       sess.run(stage_xf, feed_dict={pi: 1, x: 1, f: 2})
-      self.assertTrue(sess.run([size, isize]) == [0, 2])
+      self.assertEqual(sess.run([size, isize]), [0, 2])
 
       # Now complete key 0 with tuple entry v
       sess.run(stage_v, feed_dict={pi: 0, v: 1})
       # 1 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [1, 1])
+      self.assertEqual(sess.run([size, isize]), [1, 1])
 
       # We can now peek at 'x' and 'f' values associated with key 0
-      self.assertTrue(sess.run(peek_xf, feed_dict={pei: 0}) == {'x': 1, 'f': 2})
+      self.assertEqual(sess.run(peek_xf, feed_dict={pei: 0}), {'x': 1, 'f': 2})
       # Peek at 'v' value associated with key 0
-      self.assertTrue(sess.run(peek_v, feed_dict={pei: 0}) == {'v': 1})
+      self.assertEqual(sess.run(peek_v, feed_dict={pei: 0}), {'v': 1})
       # 1 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [1, 1])
+      self.assertEqual(sess.run([size, isize]), [1, 1])
 
       # We can now obtain 'x' and 'f' values associated with key 0
-      self.assertTrue(
-          sess.run([key_xf, get_xf], feed_dict={
-              gi: 0
-          }) == [0, {
+      self.assertEqual(
+          sess.run([key_xf, get_xf], feed_dict={gi: 0}), [0, {
               'x': 1,
               'f': 2
           }])
       # Still have 1 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [1, 1])
+      self.assertEqual(sess.run([size, isize]), [1, 1])
 
       # We can no longer get 'x' and 'f' from key 0
       with self.assertRaises(errors.InvalidArgumentError) as cm:
@@ -517,40 +515,36 @@ class MapStageTest(test.TestCase):
 
       exc_str = ("Tensor at index '0' for key '0' " 'has already been removed.')
 
-      self.assertTrue(exc_str in cm.exception.message)
+      self.assertIn(exc_str, cm.exception.message)
 
       # Obtain 'v' value associated with key 0
-      self.assertTrue(
-          sess.run([key_v, get_v], feed_dict={
-              gi: 0
-          }) == [0, {
+      self.assertEqual(
+          sess.run([key_v, get_v], feed_dict={gi: 0}), [0, {
               'v': 1
           }])
       # 0 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [0, 1])
+      self.assertEqual(sess.run([size, isize]), [0, 1])
 
       # Now complete key 1 with tuple entry v
       sess.run(stage_v, feed_dict={pi: 1, v: 1})
       # 1 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [1, 0])
+      self.assertEqual(sess.run([size, isize]), [1, 0])
 
       # Pop without key to obtain 'x' and 'f' values associated with key 1
-      self.assertTrue(sess.run([pop_key_xf, pop_xf]) == [1, {'x': 1, 'f': 2}])
+      self.assertEqual(sess.run([pop_key_xf, pop_xf]), [1, {'x': 1, 'f': 2}])
       # still 1 complete and 1 incomplete entry
-      self.assertTrue(sess.run([size, isize]) == [1, 0])
+      self.assertEqual(sess.run([size, isize]), [1, 0])
       # We can now obtain 'x' and 'f' values associated with key 1
-      self.assertTrue(
-          sess.run([pop_key_v, pop_v], feed_dict={
-              pi: 1
-          }) == [1, {
+      self.assertEqual(
+          sess.run([pop_key_v, pop_v], feed_dict={pi: 1}), [1, {
               'v': 1
           }])
       # Nothing is left
-      self.assertTrue(sess.run([size, isize]) == [0, 0])
+      self.assertEqual(sess.run([size, isize]), [0, 0])
 
   @test_util.run_deprecated_v1
   def testPartialIndexGets(self):
-    with ops.Graph().as_default() as G:
+    with ops.Graph().as_default() as g:
       with ops.device('/cpu:0'):
         x = array_ops.placeholder(dtypes.float32)
         f = array_ops.placeholder(dtypes.float32)
@@ -568,28 +562,72 @@ class MapStageTest(test.TestCase):
         size = stager.size()
         isize = stager.incomplete_size()
 
-    G.finalize()
+    g.finalize()
 
-    with self.session(graph=G) as sess:
+    with self.session(graph=g) as sess:
       # Stage complete tuple
       sess.run(stage_xvf, feed_dict={pi: 0, x: 1, f: 2, v: 3})
 
-      self.assertTrue(sess.run([size, isize]) == [1, 0])
+      self.assertEqual(sess.run([size, isize]), [1, 0])
 
       # Partial get using indices
-      self.assertTrue(
-          sess.run([key_xf, get_xf], feed_dict={
-              gi: 0
-          }) == [0, [1, 2]])
+      self.assertEqual(
+          sess.run([key_xf, get_xf], feed_dict={gi: 0}), [0, [1, 2]])
 
       # Still some of key 0 left
-      self.assertTrue(sess.run([size, isize]) == [1, 0])
+      self.assertEqual(sess.run([size, isize]), [1, 0])
 
       # Partial get of remaining index
-      self.assertTrue(sess.run([key_v, get_v], feed_dict={gi: 0}) == [0, [3]])
+      self.assertEqual(sess.run([key_v, get_v], feed_dict={gi: 0}), [0, [3]])
 
       # All gone
-      self.assertTrue(sess.run([size, isize]) == [0, 0])
+      self.assertEqual(sess.run([size, isize]), [0, 0])
+
+  @test_util.run_deprecated_v1
+  def testNonScalarKeyOrderedMap(self):
+    with ops.Graph().as_default() as g:
+      x = array_ops.placeholder(dtypes.float32)
+      v = 2. * (array_ops.zeros([128, 128]) + x)
+      t = data_flow_ops.gen_data_flow_ops.ordered_map_stage(
+          key=constant_op.constant(value=[1], shape=(1, 3), dtype=dtypes.int64),
+          indices=np.array([[6]]),
+          values=[x, v],
+          dtypes=[dtypes.int64],
+          capacity=0,
+          memory_limit=0,
+          container='container1',
+          shared_name='',
+          name=None)
+
+    g.finalize()
+
+    with self.session(graph=g) as sess:
+      with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                  'key must be an int64 scalar'):
+        sess.run(t, feed_dict={x: 1})
+
+  @test_util.run_deprecated_v1
+  def testNonScalarKeyUnorderedMap(self):
+    with ops.Graph().as_default() as g:
+      x = array_ops.placeholder(dtypes.float32)
+      v = 2. * (array_ops.zeros([128, 128]) + x)
+      t = data_flow_ops.gen_data_flow_ops.map_stage(
+          key=constant_op.constant(value=[1], shape=(1, 3), dtype=dtypes.int64),
+          indices=np.array([[6]]),
+          values=[x, v],
+          dtypes=[dtypes.int64],
+          capacity=0,
+          memory_limit=0,
+          container='container1',
+          shared_name='',
+          name=None)
+
+    g.finalize()
+
+    with self.session(graph=g) as sess:
+      with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                  'key must be an int64 scalar'):
+        sess.run(t, feed_dict={x: 1})
 
 
 if __name__ == '__main__':

programming_language: Python