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))