how to round up only k greatest elements in a Tensorflow tensor

Question

Let there be a TensorFlow tensor - for example [0.1,0.2,0.3,0.4]. I want to round up k greatest elements and round down the rest. (for example, when k =2, I wish to get [0,0,1,1]. When k =3, I wish to get [0,1,1,1].)

I want to implement this function using only TensorFlow operations. How do I achieve this?

Answer 1

Try something like this:

import tensorflow as tf

x = tf.constant([0.1,0.2,0.3,0.4])
k = 3
greatest = tf.math.top_k(x, k=k).indices 
tensor = tf.tensor_scatter_nd_update(tf.zeros_like(x), tf.reshape(greatest, (tf.shape(greatest)[0], 1)), tf.ones_like(tf.gather(x, greatest)))

k = 3:

tf.Tensor([0. 1. 1. 1.], shape=(4,), dtype=float32)

k = 2:

tf.Tensor([0. 0. 1. 1.], shape=(4,), dtype=float32)

This method does not really round, since rounding 0.3 and 0.4 to the nearest integer would results in zeros and that is not what you want. So I simply convert the highest k values in the tensor to ones and the rest to zeros, but that should be sufficient for your use case if it is still binary classification.

If you really want to round up the greatest k values, then use tf.math.ceil instead of tf.ones_like :

tensor = tf.tensor_scatter_nd_update(tf.zeros_like(x), tf.reshape(greatest, (tf.shape(greatest)[0], 1)), tf.ceil((tf.gather(x, greatest))))

Answer 2

You can use tf.math.top_k for this. The function will return the values and indices of the k-largest elements in a given tensor.

https://www.tensorflow.org/api_docs/python/tf/math/top_k

Then you can use the indices that get returned, to set the values within the tensor to a specific value.

The following solution rounds the values as mentioned in the question.

import tensorflow as tf

x = tf.constant([0.1,0.2,0.3,0.4])
k = 3

# retrieve min and max values
max_value = tf.math.ceil(tf.math.reduce_max(x))
min_value = tf.math.floor(tf.math.reduce_min(x))

# retrieve the k largest elements
k_largest = tf.math.top_k(x, k=k)

# reshape the indices, required for ‘scatter‘ function
indices = tf.reshape(k_largest.indices, (-1,1))
values = k_largest.values

# initialize update tensor with max_value
updates = max_value * tf.ones_like(values)
# initialize result with min_value
x_new = min_value * tf.ones_like(x)
# update values for k_largest indices
x_new = tf.tensor_scatter_nd_update(
    x_new, indices, updates)

print(x_new)

In case the ceil and floor operation you are asking for should be applied per element instead of applying it to the min and max value within the tensor, this would look like this:

import tensorflow as tf

x = tf.constant([0.1,0.2,0.3,0.4])
k = 3

# retrieve the k largest elements
k_largest = tf.math.top_k(x, k=k)
# reshape the indices, required for ‘scatter‘ function
indices = tf.reshape(k_largest.indices, (-1,1))

# get floored values
floored_values = tf.math.floor(x)
# get ceiled values only for top-k
ceiled_values = tf.math.ceil(k_largest.values)

# initialize result with per element floored values
x_new = floored_values
# update values for k_largest indices with per element ceiled values
x_new = tf.tensor_scatter_nd_update(
    floored_values, indices, ceiled_values)

print(x_new)