title: Map function along arbitrary dimension with pytorch? numpy?
body: <p>I wonder why I find no utility to map custom <a href="https://pytorch.org/docs/stable/index.html" rel="noreferrer">pytorch</a> or <a href="https://docs.scipy.org/doc/" rel="noreferrer">numpy</a> transformations along any dimensions of complicated tensors/arrays/matrices.<br>
  I think I remember that such a thing was available in R. With this fantasy <code>tch.map</code> utility you could do:</p>
  
  <pre><code>&gt;&gt;&gt; import torch as tch # or numpy
  
  &gt;&gt;&gt; # one torch tensor
  &gt;&gt;&gt; a = tch.tensor([0, 1, 2, 3, 4])
  &gt;&gt;&gt; # one torch function (dummy) returning 2 values
  &gt;&gt;&gt; f = lambda x: tch.tensor((x + 1, x * 2))
  &gt;&gt;&gt; # map f along dimension 0 of a, expecting 2 outputs
  &gt;&gt;&gt; res = tch.map(f, a, 0, 2) # fantasy, optimized on CPU/GPU..
  &gt;&gt;&gt; res
  tensor([[1, 0],
          [2, 2],
          [3, 4],
          [4, 6],
          [5, 8]])
  &gt;&gt;&gt; res.shape
  torch.Size([5, 2])
  
  &gt;&gt;&gt; # another tensor
  &gt;&gt;&gt; a = tch.tensor(list(range(24))).reshape(2, 3, 4).type(tch.double)
  &gt;&gt;&gt; # another function (dummy) returning 2 values
  &gt;&gt;&gt; f = lambda x: tch.tensor((tch.mean(x), tch.std(x)))
  &gt;&gt;&gt; # map f along dimension 2 of a, expecting 2 outputs
  &gt;&gt;&gt; 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]]])
  &gt;&gt;&gt; res.shape
  torch.Size([2, 3, 2])
  
  &gt;&gt;&gt; # yet another tensor
  &gt;&gt;&gt; a = tch.tensor(list(range(12))).reshape(3, 4)
  &gt;&gt;&gt; # another function (dummy) returning 2x2 values
  &gt;&gt;&gt; f = lambda x: x + tch.rand(2, 2)
  &gt;&gt;&gt; # map f along all values of a, expecting 2x2 outputs
  &gt;&gt;&gt; res = tch.map(f, a, -1, (2, 2)) # fantasy, optimized on CPU/GPU..
  &gt;&gt;&gt; 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]]]])
  &gt;&gt;&gt; res.shape
  torch.Size([3, 4, 2, 2])
  </code></pre>
  
  <p>Instead, I keep finding myself messing around with complicated <code>tch.stack</code>, <code>tch.squeeze</code>, <code>tch.reshape</code>, <code>tch.permute</code>, <em>etc.</em>, counting dimensions on my fingers not to get lost.</p>
  
  <p>Does such a utility exist and I have missed it for some reason?<br>
  Is such a utility impossible to implement for some reason?  </p>