How to Replace Values by Index in a Tensor with TensorFlow-2.0
By: Riccardo Di Sipio
Tags: TensorFlow, Machine Learning, Arrays, Numpy Array

How to Replace Values by Index in a Tensor with TensorFlow-2.0

Have you ever tried to replace only certain values of array, based on their indices?

>>> import numpy as np

>>> a = np.array( [1,2,3,4,5] )

>>> print(a) [1 2 3 4 5] >>> i = [ 0, 1, 3 ]

>>> a[i] = -1

>>> print(a) [-1 -1 3 -1 5]

A more involved use case may be, for example, to clip to a maximum all values above a certain threshold:

>>> import numpy as np

>>> a = np.array( [1,2,3,4,5] )

>>> print(a) [1 2 3 4 5] >>> a_max = 3

>>> i = np.where( a > a_max)

>>> a[i] = a_max

>>> print(a) [ 1 2 3 3 3]

It turns out that this apparently straightforward operation is not permitted in TensorFlow if the array is represented by a tensor (but it is if the array is a tf.Variable object). If you try, the most likely outcome is an error like this:

AttributeError: 'tensorflow.python.framework.ops.EagerTensor' object has no attribute 'assign'

As a matter of fact, I must admit it took a long time to figure out how to solve this problem that pops up quite often. There’s a number of questions on StackOverflow where people discussed the same issue (see this one for example). I think the easiest way to see this operation in the light of TensorFlow operations is to filter the elements that one wants to keep in one vector, create another vector to “fill the blanks” with the masked values that one wants to replace (e.g. zero, infinity, some other default value), and finally give the result as the sum of the two.

I’m giving here an example to filter the top-k elements of a vector, but the procedure is quite general and involves the following steps:

Identify the indices to remove. Typically, they are represented by a vector of boolean values such as [ True, False, False, …, True ] Convert this vector into two arrays containing the actual indices (idx_keep, idx_replace). This may be done in many ways but I believe using tf.where is the most concise method. Create a vector with the values to be replaced. In this example, the new value will be -inf, replicated N times where N = number of elements to be replaces = len(idx_replace) Create a vector with the values to be kept. I accomplished this by using the function tf.gather( a, idx_keep ) which is almost equivalent to something like a[[0,1,3]] Create sparse vectors of the two collections above. The annoying part is that indices of the sparse vectors have to be 2D even if there is only one dimension. The way I did it was to create vectors of zeros that I later “stack” with the vectors of indices. Convert the two sparse vectors into dense vectors and add them together. Return the result.

Here’s an example of the code:

A graphical representation of the workflow

Admittedly, at least for some specific cases such as the clipping described above, there exists a quicker solution based on the three-argument version of tf.where :

result = tf.where(condition, x, y)

The returned tensor contains elements of x where the condition is True and elements of y where the condition is False . So in our case the above function can be compacted to something like:

a = tf.where(tf.less_equal(a, a_max), a, a_max)

Or more explicitly:

def top_k_filtering(logits, top_k=5):

"Set values less than the k'th largest to -Inf."

kth_largest = tf.math.top_k(logits, top_k)[0][-1]

return tf.where(logits < kth_largest, -Inf, logits)

Have fun with TensorFlow!