Incorporating dim of torch.topk in tf.nn.top_k

Solution 1:

I agree with @jodag, you will have to transpose or reshape your tensor, since tf.math.top_k always works on the last dimension.

What you could also do is first get all the max values in the tensor along a certain dimension and then get the top k values from that tensor:

import tensorflow as tf
tf.random.set_seed(2)

k = 3
tensor = tf.random.uniform((2, 4, 6), maxval=10, dtype=tf.int32)
max_tensor = tf.reduce_max(tensor, axis=1)
k_max_tensor = tf.math.top_k(max_tensor, k=k, sorted=True).values

print('Original tensor --> ', tensor)
print('Max tensor --> ', max_tensor)
print('K-Max tensor --> ', k_max_tensor)
print('Unique K-Max tensor', tf.unique(tf.reshape(k_max_tensor, (tf.math.reduce_prod(tf.shape(k_max_tensor)), ))).y)

Original tensor -->  tf.Tensor(
[[[1 6 2 7 3 6]
  [7 5 1 1 0 6]
  [9 1 3 9 1 4]
  [6 0 6 2 4 0]]

 [[4 6 8 2 4 7]
  [5 0 8 2 8 9]
  [0 2 0 0 9 8]
  [9 3 8 9 0 6]]], shape=(2, 4, 6), dtype=int32)
Max tensor -->  tf.Tensor(
[[9 6 6 9 4 6]
 [9 6 8 9 9 9]], shape=(2, 6), dtype=int32)
K-Max tensor -->  tf.Tensor(
[[9 9 6]
 [9 9 9]], shape=(2, 3), dtype=int32)
Unique K-Max tensor tf.Tensor([9 6], shape=(2,), dtype=int32)