Method _testMediumTopK

tensorflow/python/kernel_tests/topk_op_test.py:142–150

(self, dtype)

    self._testLargeTopK(np.float16)

  def _testMediumTopK(self, dtype):
    b = 5
    n = 500
    k = 50
    inputs = np.random.permutation(
        np.linspace(0, 100, b * n, dtype=dtype)).reshape(b, n)
    indices = np.argsort(-inputs, axis=1)[:, :k]
    values = -np.sort(-inputs, axis=1)[:, :k]
    self._validateTopK(inputs, k, values, indices)

  def testMediumTopK(self):
    self._testMediumTopK(np.float32)