CODEWITHSUNDEEP
pythonnumpytensorflow
Felix
Felix

Reputation: 2678

Vectorised assignment to tensor

I'd like to assign multiple values to a tensor, but it seems that it's not supported at least in the way that is possible using .

a = np.zeros((4, 4))
v = np.array([0, 2, 3, 1])
r = np.arange(4)
a[r, v] = 1

>>> a
array([[1., 0., 0., 0.],
       [0., 0., 1., 0.],
       [0., 0., 0., 1.],
       [0., 1., 0., 0.]])

The above works, but the equivalent doesn't:

import tensorflow as tf

a = tf.zeros((4, 4))
v = tf.Variable([0, 2, 3, 1])
r = tf.range(4)
a[r, v].assign(1)

TypeError: Only integers, slices, ellipsis, tf.newaxis and scalar tensors are valid indices, got <tf.Tensor: shape=(4,), dtype=int32, numpy=array([0, 1, 2, 3])>

How could this be achieved? Are loops the only option? In my case the resulting array is indeed only slices of an identity matrix rearranged, so maybe that could be taken advantage of somehow.

Upvotes: 1

Views: 201

Answers (1)

Lescurel
Lescurel

Reputation: 11651

Your example, which is updating a zero tensor at some indices to a certain value is most of time achieved through tf.scatter_nd :

idx = tf.stack([r,v],axis=-1)
tf.scatter_nd(idx, updates=tf.ones(4), shape=(4,4))

For more complex cases, you can look at the following functions:

You can read more in the guide: Introduction to tensor slicing

Upvotes: 3

Related Questions