Как нарезать из многомерного массива вдоль оси в NumPy?

Question

Предположим, у меня есть NumPy ndarray x формы (n,) + higher_dims, где n - положительное целое число, а higher_dims - кортеж положительных целых чисел любой длины. То есть n - это размер первой оси, и может быть произвольно много осей.

Предположим, у меня также есть ndarray indices формы (k,) + higher_dims, где k - положительное значение int. То есть indices имеет ту же форму, что и x, за исключением, возможно, первой оси. Предположим, что каждая запись indices является целым числом между 0 и n - 1.

Я хочу создать массив y, который имеет ту же форму, что и indices, и который удовлетворяет

y[i, ...] = x[indices[i, ...], ...]

для каждого i между 0 и n - 1. Здесь ... обозначает произвольную комбинацию индексов для оставшихся осей, а не объект эллипса .

Например, вот как я мог бы создать y, если x является трехмерным с использованием for-loop:

import numpy as np

x = np.arange(24).reshape((4, 2, 3))
print('x =', x, sep='\n')

indices = np.asarray([[[1, 0, 1], [2, 1, 2]], [[3, 1, 2], [0, 0, 1]]])
print('indices =', indices, sep='\n')

y = np.empty(indices.shape, dtype=x.dtype)
for i in range(indices.shape[0]):
    for j in range(indices.shape[1]):
        for k in range(indices.shape[2]):
            y[i, j, k] = x[indices[i, j, k], j, k]  # Defining property of y
print('y =', y, sep='\n')

Выход:

x =
[[[ 0  1  2]
  [ 3  4  5]]
 [[ 6  7  8]
  [ 9 10 11]]
 [[12 13 14]
  [15 16 17]]
 [[18 19 20]
  [21 22 23]]]
indices =
[[[1 0 1]
  [2 1 2]]
 [[3 1 2]
  [0 0 1]]]
y =
[[[ 6  1  8]
  [15 10 17]]
 [[18  7 14]
  [ 3  4 11]]]

Я ищу функцию или трюк индексирования для достижения такого поведения в целом (для ndarrays произвольного измерения), без петель Python, если это возможно .