How to fill a tensor of values based on tensor of indices in tensorflow?

user_5561472 0
Viewed 285

I need to extract values from tensor based on the indices tensor.

My code is as follows:

arr = tf.constant([10, 11, 12]) # array of values
inds = tf.constant([0, 1, 2])   # indices
res = tf.map_fn(fn=lambda t: arr[t], elems=inds)

It works slowly. Is there more efficient way ?

tensorflow python
1 Answers

You can use tf.gather method

  arr = tf.constant([10, 11, 12]) # array of values
  inds = tf.constant([0, 2]) 

  r = tf.gather(arr , inds)#<tf.Tensor: shape=(2,), dtype=int32, numpy=array([10, 12])>

If you have a multi-dimensional tensor, The tf.gather has an "axis" param to specify the dimension where you check the indices :

arr = tf.constant([[10, 11, 12] ,[1, 2, 3]]) # shape(2,3)
inds = tf.constant([0, 1]) 
# axis == 1
r = tf.gather(arr , inds , axis = 1)#<tf.Tensor: shape=(2, 2), dtype=int32, numpy=array([[10, 11],[ 1,  2]])>

 
# axis == 0
 r = tf.gather(arr , inds , axis = 0) #<tf.Tensor: shape=(2, 3), dtype=int32, numpy=array([[10, 11, 12], [ 1,  2,  3]])>
user_52355280
Related