Repository: pytorch
Issue: Pipe method for torch tensors
Labels: module: dataloader,triaged

## 🚀 Feature

Pandas DataFrames include a `pipe()` method, as documented [here](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.pipe.html). In short, this method allows a DataFrame to be processed by any user-defined function which accepts and returns a DataFrame, without dropping out of method chaining.
It would be great to add this functionality to `torch.Tensor`.

## Motivation

The `pipe()` method allows one to stay within method chaining, without having to break connected operations into distinct lines.

### Example 1

In the code below I compute the FFT of a signal and then perform some additional operations on the result.

```python
def _zero_below_mean(x: torch.Tensor) -> torch.Tensor:
    x[x < x.mean()] = 0.0
    return x

x = torch.randn(1, 1024)

result = (
    torch.fft.fft(x)
    .abs()
    .add(1e-8)
    .log()
)
result = _zero_below_mean(result).mean()
```

Here, I've had to break the code into two pieces: before `_zero_below_mean()` and after it.

Conversely, with a `pipe()` method it would be possible to write this code much more succinctly: 

```python
result = (
    torch.fft.fft(x)
    .abs()
    .add(1e-8)
    .log()
    .pipe(_zero_below_mean)
    .mean()
)
```

This is, of course, a completely contrived example. 😄 However, I hope it is illustrative nonetheless.

### Example 2

In this example I define an "Up Convolution" layer, that simply invokes part of the functional API to interpolate between convolutions.

Without a pipe method:

```python
from torch import nn


class UpConv2d(nn.Module):
    def __init__(self, in_channels: int, out_channels: int) -> None:
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels

        self.conv0 = nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=1)
        self.conv1 = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = self.conv0(x)
        y = F.interpolate(y, scale_factor=2)
        y = self.conv1(y)
        return y
```

With a pipe method:

```python
class UpConv2d(nn.Module):
    def __init__(self, in_channels: int, out_channels: int) -> None:
        super().__init__()
        ...
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return (
            self.conv0(x)
            .pipe(F.interpolate, scale_factor=2)
            .pipe(self.conv1)
        )
```

In my experience, the former approach can also increase the chance of introducing a subtle and highly consequential bug into the code, such as this:

```python
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        y = self.conv0(x)
        y = F.interpolate(x, scale_factor=2)
        y = self.conv1(y)
        return y
```

Above, the first convolution is not truly utilized, and the mistake can be very hard to notice.
Such an error is highly improbable in the second approach.

## Pitch

Implementing a `pipe()` operator _seems_ to be relatively simple, as shown below. (Granted, there may be complexities I am unaware of.)

```python
import torch
from typing import Any, Callable


class Tensor(torch.Tensor):
    def pipe(
        self,
        func: Callable[[torch.Tensor], torch.Tensor],  # typed this way for brevity.
        *args: Any,
        **kwargs: Any,
    ) -> torch.Tensor:
        return func(self, *args, **kwargs)
```

## Alternatives

The alternative here is to not implement a `pipe()` method on tensors. Such an argument could be motivated by concerns about readability, performance, etc. (I am, of course, partial to the view that `pipe()` would improve readability and am not aware of any performance issues, but would be happy to hear what others think.)

## Additional context

\-


cc @SsnL @VitalyFedyunin @ejguan @NivekT