Карта функции вдоль произвольного измерения с pytorch?NumPy?

Question

Интересно, почему я не нахожу утилиту для отображения пользовательских pytorch или numpy преобразований по любым измерениям сложных тензоров / массивов / матриц.
Мне кажется, я помню, что такоеэта штука была доступна в R. С помощью этой фантастической tch.map утилиты, которую вы могли бы сделать:

>>> import torch as tch # or numpy

>>> # one torch tensor
>>> a = tch.tensor([0, 1, 2, 3, 4])
>>> # one torch function (dummy) returning 2 values
>>> f = lambda x: tch.tensor((x + 1, x * 2))
>>> # map f along dimension 0 of a, expecting 2 outputs
>>> res = tch.map(f, a, 0, 2) # fantasy, optimized on CPU/GPU..
>>> res
tensor([[1, 0],
        [2, 2],
        [3, 4],
        [4, 6],
        [5, 8]])
>>> res.shape
torch.Size([5, 2])

>>> # another tensor
>>> a = tch.tensor(list(range(24))).reshape(2, 3, 4).type(tch.double)
>>> # another function (dummy) returning 2 values
>>> f = lambda x: tch.tensor((tch.mean(x), tch.std(x)))
>>> # map f along dimension 2 of a, expecting 2 outputs
>>> res = tch.map(f, a, 2, 2) # fantasy, optimized on CPU/GPU..
tensor([[[ 1.5000,  1.2910],
         [ 5.5000,  1.2910],
         [ 9.5000,  1.2910]],

        [[13.5000,  1.2910],
         [17.5000,  1.2910],
         [21.5000,  1.2910]]])
>>> res.shape
torch.Size([2, 3, 2])

>>> # yet another tensor
>>> a = tch.tensor(list(range(12))).reshape(3, 4)
>>> # another function (dummy) returning 2x2 values
>>> f = lambda x: x + tch.rand(2, 2)
>>> # map f along all values of a, expecting 2x2 outputs
>>> res = tch.map(f, a, -1, (2, 2)) # fantasy, optimized on CPU/GPU..
>>> print(res)
tensor([[[[ 0.4827,  0.3043],
          [ 0.8619,  0.0505]],

         [[ 1.4670,  1.5715],
          [ 1.1270,  1.7752]],

         [[ 2.9364,  2.0268],
          [ 2.2420,  2.1239]],

         [[ 3.9343,  3.6059],
          [ 3.3736,  3.5178]]],


        [[[ 4.2063,  4.9981],
          [ 4.3817,  4.4109]],

         [[ 5.3864,  5.3826],
          [ 5.3614,  5.1666]],

         [[ 6.6926,  6.2469],
          [ 6.7888,  6.6803]],

         [[ 7.2493,  7.5727],
          [ 7.6129,  7.1039]]],


        [[[ 8.3171,  8.9037],
          [ 8.0520,  8.9587]],

         [[ 9.5006,  9.1297],
          [ 9.2620,  9.8371]],

         [[10.4955, 10.5853],
          [10.9939, 10.0271]],

         [[11.3905, 11.9326],
          [11.9376, 11.6408]]]])
>>> res.shape
torch.Size([3, 4, 2, 2])

Вместо этого я продолжаю баловаться со сложными tch.stack, tch.squeeze, tch.reshape, tch.permute, и т. Д. , считая размеры на моих пальцах, чтобы не потеряться.

Существует ли такая утилита, и я почему-то пропустил ее?
Такая утилита невозможна дляреализовать по какой-то причине?