How to use an array to index the last dim of a second array

Assume that X.shape = (h, w, d) and that Y.shape = (h, w) contains values (indices) in range(d).

How can I take the elements in X (from the last dim) using the indices in Y?

That is, I would like to do something like X[Y] that will return an h x w array where Y was used as the indices in the last dim d.

I think you are looking for numpy.take_along_axis:

import numpy as np
rng = np.random.default_rng()

h, w, d = 2, 3, 4
x = rng.random((h, w, d))
y = rng.integers(0, d, (h, w))

x, y, np.take_along_axis(x, y[..., None], axis=-1)[..., 0]

Output:

(array([[[0.51705108, 0.68891581, 0.84475703, 0.77938839],
         [0.02115493, 0.47689898, 0.19786926, 0.73959225],
         [0.40821923, 0.0119006 , 0.89595898, 0.81798467]],
 
        [[0.60350791, 0.11501983, 0.15932539, 0.35923036],
         [0.27939872, 0.13691148, 0.47528086, 0.71320657],
         [0.98294212, 0.75039413, 0.06087527, 0.68233282]]]),
 array([[0, 2, 0],
        [3, 0, 0]]),
 array([[0.51705108, 0.19786926, 0.40821923],
        [0.35923036, 0.27939872, 0.98294212]]))

Comparison with Ulises' answer:

assert np.array_equal(np.take_along_axis(x, y[..., None], axis=-1)[..., 0], x[ np.eye(d)[y].astype(np.bool)].reshape(y.shape))
%timeit np.take_along_axis(x, y[..., None], axis=-1)[..., 0]
%timeit x[ np.eye(d)[y].astype(np.bool)].reshape(y.shape)

Output:

9.42 µs ± 91.2 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
55.8 µs ± 463 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

Edit

The shortest way to select elements is swaping dimensions in x then use choose to select y values:

np.choose(y, x.transpose(2,0,1))

Old solution

suposing you have x, and y, you could create a categorical array using y and use it to select on x:

example:

#create arrays
h,w,d = 3,3,4
x = np.arange(h*w*d).reshape(h,w,d)
y = np.random.randint(d,size=(h,w))

#categorical y
cat_y =  np.eye(d)[y]

# select elements from x
x_sel = x[cat_y.astype(bool)]

#reshape to original form
x_out = x_sel.reshape(y.shape)

or in one line:

xx = x[ np.eye(d)[y].astype(np.bool)].reshape(y.shape)

X[np.arange(h)[:,None], np.arange(w), Y]

The idea is to index the first two dim with arrays that broadcast with Y. (h,1), (w,), (h,w) That's the general principle for all advanced indexing. That's we had to do before they added take_along.

ix_, meshgrid and ogrid could also be used to make these arrays.