tf.gather_nd to get values of 2d array from 1d array of indices

user_11656914 0
Viewed 262

I have a 2D array of values and a 1D array of indices. I want to get the value of each index at each element of the 2D array using tf.gather_nd or some other numpy or tensorflow command like so:

2D_array = np.arange(27).reshape((9, 3))
array([[ 0,  1,  2],
       [ 3,  4,  5],
       [ 6,  7,  8],
       [ 9, 10, 11],
       [12, 13, 14],
       [15, 16, 17],
       [18, 19, 20],
       [21, 22, 23],
       [24, 25, 26]])

1D_array = np.random.randint(low=0, high=3, size=9)
array([2, 2, 0, 2, 0, 1, 0, 1, 0])

output
[2, 5, 6, 11, 12, 16, 18, 22, 24]

What I currently have is this mess:

output = tf.gather_nd(a, list(zip(np.arange(len(b)), b)))
<tf.Tensor: shape=(9,), dtype=int32, numpy=array([ 2,  5,  6, 11, 12, 16, 18, 22, 24])>

A cleaner solution that avoids using zip and list would be greatly appreciated, as I don't think I'm currently utilizing tf.gather_nd completely correctly.

tensorflow numpy arrays python
1 Answers

The 1D_array in your code should be made a 2d array with

inds = np.random.randint(low=0, high=3, size=9)
inds = np.dstack((np.arange(len(inds)),inds))

with this 2d array of indices, you can directly apply tf.gather_nd

import tensorflow as tf 
import numpy as np

tmp = np.arange(27).reshape((9, 3))
inds = np.random.randint(low=0, high=3, size=9)
inds = np.dstack((np.arange(len(inds)),inds))

output = tf.gather_nd(tmp,inds)

with tf.Session() as sess:
    output_val = sess.run(output)
user_83668050
Related