# Apply a function along an axis

**URL:** <https://discuss.pytorch.org/t/apply-a-function-along-an-axis/130440>\
**Category:** Uncategorized\
**Created:** [August 27, 2021, 9:57am UTC](https://discuss.pytorch.org/t/apply-a-function-along-an-axis/130440 "2021-08-27T09:57:43Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![thomas](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/t/e36b37/32.png) [@thomas](https://discuss.pytorch.org/u/thomas)\
**Post date:** [August 27, 2021, 9:57am UTC](https://discuss.pytorch.org/t/apply-a-function-along-an-axis/130440/1 "2021-08-27T09:57:43Z")

</div>

Hello everyone!

I have an input of this shape:

```auto
(num_samples, num_symbols, num_features, num_observations)

```

I would like to feed through this input through a neural network.

The neural network as `(num_features, num_observations)` shape input and `(num_outputs)` outputs, giving me `(num_samples, num_symbols, num_outputs)` when I apply along

```auto
def apply_along_axis(function, x, axis: int = 0):
    return torch.stack([
        function(x_i) for x_i in torch.unbind(x, dim=axis)
    ], dim=axis)

```

(from [Pytorch equivalent of numpy.apply\_along\_axis](https://discuss.pytorch.org/t/pytorch-equivalent-of-numpy-apply-along-axis/3422) )

This is an operation that is made very often in my network, and I just discover that 42% of the time is spent in the apply\_along\_axis.

I wanted to know if there is better way to do this ? (the for loop seems to be quite slow)

Is there a possibility to create identical versions of the same “function” in the network and update all the weights automatically when grad is invoked ?

(I maybe have a similar question with more complex situation if there is good suggestion related to this issue)

Thanks!

---

<div class="post-metadata">

**Author:** ![KFrank](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/k/ecb155/32.png) [@KFrank](https://discuss.pytorch.org/u/KFrank)\
**Post date:** [August 27, 2021, 4:33pm UTC](https://discuss.pytorch.org/t/apply-a-function-along-an-axis/130440/2 "2021-08-27T16:33:25Z")

</div>

Hi Thomas!

> [@thomas](#):
>
> I have an input of this shape:
> 
> ```jsx
> (num_samples, num_symbols, num_features, num_observations)
> 
> ```
> 
> I would like to feed through this input through a neural network.
> 
> …
> 
> ```jsx
> def apply_along_axis(function, x, axis: int = 0):
> return torch.stack([
> function(x_i) for x_i in torch.unbind(x, dim=axis)
> ], dim=axis)
> 
> ```
> 
> I wanted to know if there is better way to do this ?

I do not know of any functionality built into pytorch similar to your  
`apply_along_axis()`. And even if there were, it would still impose  
a performance penalty, as it would still be breaking up what might  
have been a single, larger tensor operation into many smaller, axis-wise  
tensor operations (absent some kind of hypothetical JIT compilation).

As a general rule, if you find yourself looping over a tensor, you should  
see if you can recast your computation into pure (that is, loop-free)  
tensor operations. Sometimes you can and sometimes you can’t.

Note that you can _sometimes_ realize a net performance gain by getting  
rid of loops, _even if_ your loop-free approach is of higher computational  
complexity.

Here is a simplistic, contrived example that illustrates replacing  
`apply_along_axis()` with a single pytorch tensor operation:

```plaintext
>>> import torch
>>> torch. __version__'1.9.0'
>>> _ = torch.manual_seed (2021)
>>> def apply_along_axis(function, x, axis: int = 0):
... return torch.stack([
... function(x_i) for x_i in torch.unbind(x, dim=axis)
... ], dim=axis)
...
>>> def my_fn (x):
... return torch.softmax (x, 0)
...
>>> t = torch.randn (2, 3)
>>> t
tensor([[2.2871, 0.6413, -0.8615],
        [-0.3649, -0.6931, 0.9023]])
>>> apply_along_axis (my_fn, t, 0)
tensor([[0.8092, 0.1561, 0.0347],
        [0.1897, 0.1366, 0.6737]])
>>> torch.softmax (t, 1)
tensor([[0.8092, 0.1561, 0.0347],
        [0.1897, 0.1366, 0.6737]])
>>> apply_along_axis (my_fn, t, 1)
tensor([[0.9341, 0.7916, 0.1463],
        [0.0659, 0.2084, 0.8537]])
>>> torch.softmax (t, 0)
tensor([[0.9341, 0.7916, 0.1463],
        [0.0659, 0.2084, 0.8537]])

```

You suggest that your use case involves having the `function` you apply  
be an entire neural-network model.

Although various model layers do have constraints on the shapes they  
expect, the basic building blocks accept (and sometimes require) an  
arbitrary batch dimension. This suggests that reworking your model so  
that you don’t need `apply_along_axis` could be plausible.

Two building-block examples: `Linear` accepts an arbitrary number  
of leading “batch” dimensions, so that’s likely to be easy. On the other  
hand, `Conv2d` requires a tensor of exactly four dimensions, but it’s  
leading dimension is an arbitrary batch dimension, so you can use  
`view()` (or `reshape()`) to repackaged multiple “batch” dimensions  
into a single batch dimension. Thus:

```plaintext
>>> import torch
>>> torch. __version__'1.9.0'
>>> torch.nn.Linear (2, 3) (torch.randn (7, 7, 5, 2)).shape
torch.Size([7, 7, 5, 3])
>>> torch.nn.Conv2d (2, 3, 3, padding = 1) (torch.randn (7, 7, 5, 2, 11, 11).view (7*7*5, 2, 11, 11)).view (7, 7, 5, 3, 11, 11).shape
torch.Size([7, 7, 5, 3, 11, 11])

```

Whether or not you would be able to push these kinds of techniques  
through an entire model will depend on the model’s details, but there  
are certainly some realistic models where you could.

Best.

K. Frank