Map function along arbitrary dimension with pytorch? numpy?

Question

I wonder why I find no utility to map custom pytorch or numpy transformations along any dimensions of complicated tensors/arrays/matrices.
I think I remember that such a thing was available in R. With this fantasy tch.map utility you could do:

>>> 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])

Instead, I keep finding myself messing around with complicated tch.stack , tch.squeeze , tch.reshape , tch.permute , etc. , counting dimensions on my fingers not to get lost.

Does such a utility exist and I have missed it for some reason?
Is such a utility impossible to implement for some reason?