概要

Python で多重配列 numpy.ndarray もしくは torch.Tensor からブロック多重配列をスライスするときは関数 numpy.ix_ を使う。

詳細

Julia で

# Julia
# 多重配列を作る
(i, j, k) = (2, 3, 4)
x = reshape(1:i * j * k, i, j, k)
# 多重配列を多重添字でスライスする
(is, js, ks) = ([2, 1], [1, 3], 2:3)
x[is, js, ks] |> println

みたいにして多重配列を多重添字でスライスするのを Python でどうやるのか調べてもなかなか情報にたどり着かなかったのでメモ。numpy.ix_ が所望のものである。 NumPy や PyTorch はデフォルトだと Julia の Array とは多重配列の次元の並べ方が逆であることに注意。

# Python + NumPy
import numpy
# 多重配列を作る
i, j, k = 2, 3, 4
x = (numpy.arange(k * j * i) + 1).reshape(k, j, i)
# 多重配列を多重添字でスライスする
# (Python では is は予約語かなんかのようなので is_ とする)
is_, js, ks = numpy.array([2, 1]) - 1, numpy.array([1, 3]) - 1, numpy.arange(2 - 1, 3)
print(x[numpy.ix_(ks, js, is_)])

# Python + PyTorch
import torch
import numpy
# 多重配列を作る
i, j, k = 2, 3, 4
x = (torch.arange(k * j * i) + 1).reshape(k, j, i)
# 多重配列を多重添字でスライスする
# (Python では is は予約語かなんかのようなので is_ とする)
is_, js, ks = torch.tensor([2, 1]) - 1, torch.tensor([1, 3]) - 1, torch.arange(2 - 1, 3)
print(x[numpy.ix_(ks, js, is_)])