Digamos que eu tenha uma série de permutações permque poderiam ser assim:
perm = np.array([[0, 1, 2], [1, 2, 0], [0, 2, 1], [2, 1, 0]])
Se eu quiser aplicá-lo a um eixo, posso escrever algo como:
v = np.arange(9).reshape(3, 3)
print(v[perm])
Saída:
array([[[0, 1, 2],
[3, 4, 5],
[6, 7, 8]],
[[3, 4, 5],
[6, 7, 8],
[0, 1, 2]],
[[0, 1, 2],
[6, 7, 8],
[3, 4, 5]],
[[6, 7, 8],
[3, 4, 5],
[0, 1, 2]]])
Agora gostaria de aplicá-lo a dois eixos ao mesmo tempo. Eu descobri que posso fazer isso através de:
np.array([v[tuple(np.meshgrid(p, p, indexing="ij"))] for p in perm])
Mas acho isso bastante ineficiente, porque precisa criar uma grade de malha e também requer um loop for. Criei um array pequeno neste exemplo, mas na realidade tenho arrays muito maiores com muitas permutações, então adoraria ter algo que fosse tão rápido e simples quanto a versão de um eixo.