Array of indexes for each element alongs the first dimension in a 2D array (numpy., tensorflow)

Question

indexes = np.array([[0,1,3],[1,2,4 ]])
data = np.random.rand(2,5)

Now, i would like an array of shape (2,3), where

result[0] = data[0,indexes[0]]
result[1] = data[1,indexes[1]]

What would be the proper way to achieve this? A numpy way that yould generalize to bigger arrays (perhaps even higher dimensional).

Please note the difference to questions like this, where the array of indexes contains tuples. This is not what I am asking.

Edit

A more general formulation of the question would be:

data.shape == (s0, s1, .., sn)
indexes.shape == (s0, s1, ..., sn-1, K)
so, they have all dimension but the last equal

Than

result[i, j, ..., k] = data[i, j,...,k, indexes[i, j, ..., k]]

where

len([i, j, ..., k]) == len(data)-1 == len(indexes) - 1

score 1 · Accepted Answer · answered Dec 06 '18 at 14:16

Here are NumPy and TensorFlow solutions:

import numpy as np
import tensorflow as tf

def gather_index_np(data, index):
    data = np.asarray(data)
    index = np.asarray(index)
    # Make open grid of all but last dimension indices
    grid = np.ogrid[tuple(slice(s) for s in index.shape[:-1])]
    # Add extra dimension in grid
    grid = [g[..., np.newaxis] for g in grid]
    # Complete index
    index_full = tuple(grid + [index])
    # Index data to get result
    result = data[index_full]
    return result

def gather_index_tf(data, index):
    data = tf.convert_to_tensor(data)
    index = tf.convert_to_tensor(index)
    index_shape = tf.shape(index)
    d = index.shape.ndims
    # Make grid of all dimension indices
    grid = tf.meshgrid(*(tf.range(index_shape[i]) for i in range(d)), indexing='ij')
    # Complete index
    index_full = tf.stack(grid[:-1] + [index], axis=-1)
    # Index data to get result
    result = tf.gather_nd(data, index_full)
    return result

Example:

import numpy as np
import tensorflow as tf

data = np.arange(10).reshape((2, 5))
index = np.array([[0, 1, 3], [1, 2, 4]])
print(gather_index_np(data, index))
# [[0 1 3]
#  [6 7 9]]
with tf.Session() as sess:
    print(sess.run(gather_index_tf(data, index)))
# [[0 1 3]
#  [6 7 9]]

score 1 · Answer 2 · answered Dec 06 '18 at 18:15

numpy has take_along_axis which does what you describe plus it also lets you choose the axis.

Example:

>>> a = np.arange(24).reshape(2,3,4)
>>> i = np.random.randint(0,4,(2,3,5))
>>> i
array([[[3, 3, 0, 1, 3],
        [3, 1, 0, 3, 3],
        [3, 2, 0, 3, 3]],

       [[2, 3, 0, 0, 0],
        [1, 1, 3, 1, 2],
        [1, 3, 0, 0, 2]]])

>>> np.take_along_axis(a, i, -1)
array([[[ 3,  3,  0,  1,  3],
        [ 7,  5,  4,  7,  7],
        [11, 10,  8, 11, 11]],

       [[14, 15, 12, 12, 12],
        [17, 17, 19, 17, 18],
        [21, 23, 20, 20, 22]]])

TensorFlow has an experimental version of that NumPy function ```tf.experimental.numpy.take_along_axis``` — Joe, Aug 05 '21 at 20:33

Array of indexes for each element alongs the first dimension in a 2D array (numpy., tensorflow)

Edit

2 Answers2