矢量化numpy中循环的五元组，具有非平凡的索引

Question

我有一系列嵌套for循环，用于填充秩3数组，如下所示：

S = ...  # Some number
N = ...  # Some other number
an_array = np.zeros((S, N * N, N * N))
for s in range(S):
    for i in range(N):
        for a in range(N):
            for j in range(N):
                for b in range(N):
                    an_array[s, i * N + a, j * N + b] = target[s, i, a, j, b]

这相当于在数组最右边的两个维度的基N中有索引，其中数组的一个元素被索引为an_array[s, ia, jb]，只是索引在基10中转换以使用numpy

如何以矢量化方式填充数组而不诉诸for循环

谢谢