def testTopKthIterativeValue(self):
x = [1.0, 2.0, 3.0, 4.0]
y = common_layers.top_kth_iterative(tf.constant(x, dtype=tf.float32),
3)
actual = self.evaluate(y)
self.assertEqual(int(actual[0]), 2.0)
def testTopKthIterativeShape(self):
x = np.random.rand(5, 2, 1, 12)
y = common_layers.top_kth_iterative(tf.constant(x, dtype=tf.float32),
3)
actual = self.evaluate(y)
self.assertEqual(actual.shape, (5, 2, 1, 1))
