How to index a tensor and change the value

user_13818399 0
Viewed 854

I am working on a algorithm with tensorflow. Following is the NumPy version of the wanted code:

x = [1,2,3,4,5,6,7,8,9,10]
sets = {1,5,7}
y = [0,0,0,0,0,0,0,0,0,0]

for i in range(10):
    if i in sets:
        y[i] = x[i]

which gets the result:

y = [0,2,0,0,0,6,0,8,0,0]

How to implement this in tensorflow? Is there any way to to implement this in tensorflow using the same logic, not by converting NumPy arrays to tensor after the calculation, but instead doing all the operation using tensor (e.g. using tensor to index the tensor, and assign its value by x(tensor) if the index is in the set).

tensorflow python
2 Answers

You can do that in TensorFlow like this:

import tensorflow as tf
x = tf.constant([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
sets = tf.constant([1, 5, 7])
y = tf.constant([0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
y2 = tf.tensor_scatter_nd_update(y, tf.expand_dims(sets, 1), tf.gather(x, sets))
print(y2.numpy())
# [0 2 0 0 0 6 0 8 0 0]

If your y is always made of zeros (that is, you only want to "fill" the positions given by sets), then you can just do:

import tensorflow as tf
x = tf.constant([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
sets = tf.constant([1, 5, 7])
y2 = tf.scatter_nd(tf.expand_dims(sets, 1), tf.gather(x, sets), tf.shape(x))
print(y2.numpy())
# [0 2 0 0 0 6 0 8 0 0]

As an alternative method, you can do it with masking too, although I don't think it should be faster:

import tensorflow as tf
x = tf.constant([1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
sets = tf.constant([1, 5, 7])
y = tf.constant([0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
s = tf.shape(x, out_type=sets.dtype)
mask = tf.math.reduce_any(tf.equal(tf.range(s[0]), tf.expand_dims(sets, 1)), 0)
y2 = tf.where(mask, x, y)
# Or if y is always zeros:
#y2 = x * tf.dtypes.cast(mask, x.dtype)
print(y2.numpy())
# [0 2 0 0 0 6 0 8 0 0]
user_17827920

I don't clearly understood your question. But here's what i think you need. You can convert numpy array to tensor using

test_array = np.array([1,2,3,4])
random1_array = tf.convert_to_tensor(test_array)

Here's an alternative way if you don't want to use numpy and convert it to tensor.

test_array = tf.Variable([1,2,3,4])
test_array_2 = tf.Variable([5,6,7,8])

now you can do all sort of calculations like add, substract, etc. like you do with numpy arrays.

user_137651720
Related