Method to get k smallest elements from a tensor

[digital-idiot]

tf.math.top_k only returns k largest elements from the input tensor along a specified axis. However, there is no direct way to obtain k smallest elements from a tensor. The equivalent method topk in pytorch has an bool flag largest to toggle the output from k largest elements to k smallest elements and vice versa. It would be useful to have similar option (or an alternative method like bottom_k) for this purpose so that hacky solution like the following can be avoided.

tf.negative(tf.math.top_k(tf.negative(A)))

Finding k minimum values can be potentially very useful in many cases especially for writing custom ops. This does not seem to be too hard to implement or drastically change the API.

System information

• TensorFlow version: 2.3

Related discussion in Stackoverflow.