Repository: pytorch
Issue: [JIT] Infinite RecursionError with self-referential models (also affects `__repr__`)!!
Labels: oncall: jit

### 🐛 Describe the bug

So, this may be a bit unorthodox, but I am working with invertible models, i.e. models for (left/right) invertible functions. These are `nn.Module`s that not only have a `.forward`, but also a `.inverse`. A popular example would be the [invertible Residual Network](https://proceedings.mlr.press/v97/behrmann19a.html).

For such models, it would be nice to have an `.invert()` function or  `.inverted` property/submodule, that returns the inverse function as an `nn.Module`, which has the roles of `.forward` and `.inverse` swapped. Ideally, I would even want to implement this with the `operator.__inv__` method, such that 

```python
assert ~(~model) is model
assert model.inverted.inverted is model
```
holds true. This is particularly desirable when using such modules as encoder/decoder components in a larger network.

Now, I already figured out a way to create such models without `jit` (full MWE at the end), by simply putting the appropriate references. However, `str(model)` / `repr(model)` fails spectacularly due to infinite recursion. Apart from that these models can be used just fine.

<details><summary>Full MWE</summary>

```python
from typing import Final, Optional

import torch
from torch import Tensor, jit, nn


# @autojit
class ConcatProjection(nn.Module):
    r"""Maps `z = [x,w] ⟼ x`."""

    # Constants
    input_size: Final[int]
    r"""CONST: The dimensionality of the inputs."""
    hidden_size: Final[int]
    r"""CONST: The dimensionality of the outputs."""
    pad_size: Final[int]
    r"""CONST: The size of the padding."""

    # Parameters
    padding: Tensor
    r"""PARAM: The padding vector."""

    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        padding: Optional[Tensor] = None,
        inverted: Optional[nn.Module] = None,
    ) -> None:
        super().__init__()

        if not input_size >= hidden_size:
            raise ValueError(
                f"ConcatProjection requires {input_size=} ≥ {hidden_size=}!"
            )

        self.input_size = input_size
        self.hidden_size = hidden_size
        self.pad_size = input_size - hidden_size

        if padding is None:
            padding = nn.Parameter(torch.randn(self.pad_size))
        elif not isinstance(padding, nn.Parameter):
            padding = nn.Parameter(padding)

        self.register_parameter("padding", padding)

        if inverted is None:
            inverted = ConcatEmbedding(
                input_size=self.hidden_size,
                hidden_size=self.input_size,
                padding=self.padding,
                inverted=self,
            )
        self.inverted = inverted

    @jit.export
    def __invert__(self):
        return self.inverted

    @jit.export
    def forward(self, Z: Tensor) -> Tensor:
        r"""Signature: `[..., d+e] ⟶ [..., d]`."""
        return Z[..., : self.hidden_size]

    @jit.export
    def inverse(self, X: Tensor) -> Tensor:
        r"""Signature: `[..., d] ⟶ [..., d+e]`."""
        shape = list(X.shape[:-1]) + [self.pad_size]
        return torch.cat([X, self.padding.expand(shape)], dim=-1)


# @autojit
class ConcatEmbedding(nn.Module):
    r"""Maps `x ⟼ [x,w]`."""

    # Constants|
    input_size: Final[int]
    r"""CONST: The dimensionality of the inputs."""
    hidden_size: Final[int]
    r"""CONST: The dimensionality of the outputs."""
    pad_size: Final[int]
    r"""CONST: The size of the padding."""

    # Parameters
    padding: Tensor
    r"""PARAM: The padding vector."""

    def __init__(
        self,
        *,
        input_size: int,
        hidden_size: int,
        padding: Optional[Tensor] = None,
        inverted: Optional[nn.Module] = None,
    ) -> None:
        super().__init__()

        if not input_size <= hidden_size:
            raise ValueError(
                f"ConcatEmbedding requires {input_size=} ≤ {hidden_size=}!"
            )

        self.input_size = input_size
        self.hidden_size = hidden_size
        self.pad_size = hidden_size - input_size

        padding = None
        if padding is None:
            padding = nn.Parameter(torch.randn(self.pad_size))
        elif not isinstance(padding, nn.Parameter):
            padding = nn.Parameter(padding)

        self.register_parameter("padding", padding)

        if inverted is None:
            inverted = ConcatProjection(
                input_size=self.hidden_size,
                hidden_size=self.input_size,
                padding=self.padding,
                inverted=self,
            )
        self.inverted = inverted

    @jit.export
    def __invert__(self):
        return self.inverted

    @jit.export
    def forward(self, X: Tensor) -> Tensor:
        r"""Signature: `[..., d] ⟶ [..., d+e]`."""
        shape = list(X.shape[:-1]) + [self.pad_size]
        return torch.cat([X, self.padding.expand(shape)], dim=-1)

    @jit.export
    def inverse(self, Z: Tensor) -> Tensor:
        r"""Signature: `[..., d+e] ⟶ [..., d]`."""
        return Z[..., : self.input_size]

model = ConcatEmbedding(input_size=4, hidden_size=16)
assert model is ~(~model)

model = ConcatProjection(input_size=16, hidden_size=4)
assert model is ~(~model)

print(model)  # RecursionError: maximum recursion depth exceeded
jit.script(model)  # RecursionError: maximum recursion depth exceeded while getting the str of an object
```
</details>

### Versions


<details><summary>Environment</summary>

```
Collecting environment information...
PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 21.10 (x86_64)
GCC version: (Ubuntu 11.2.0-7ubuntu2) 11.2.0
Clang version: 13.0.0-2
CMake version: Could not collect
Libc version: glibc-2.34

Python version: 3.9.12 | packaged by conda-forge | (main, Mar 24 2022, 23:25:59)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.13.0-40-generic-x86_64-with-glibc2.34
Is CUDA available: False
CUDA runtime version: 11.5.119
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.4.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.4.0
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.1.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.1
/usr/local/cuda-11.3/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.1
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn.so.8.2.4
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.2.4
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.2.4
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.2.4
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.2.4
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.2.4
/usr/local/cuda-11.4/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.2.4
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] mypy==0.942
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.21.6
[pip3] numpydoc==1.3.1
[pip3] pytorch-forecasting==0.10.1
[pip3] pytorch-lightning==1.6.2
[pip3] torch==1.11.0
[pip3] torch-tb-profiler==0.4.0
[pip3] torchaudio==0.11.0
[pip3] torchdata==0.4.0a0+b739701
[pip3] torchdiffeq==0.2.2
[pip3] torchinfo==1.6.5
[pip3] torchmetrics==0.8.1
[pip3] torchtext==0.12.0
[pip3] torchvision==0.12.0
[conda] blas                      2.114                       mkl    conda-forge
[conda] blas-devel                3.9.0            14_linux64_mkl    conda-forge
[conda] cudatoolkit               11.3.1               ha36c431_9    nvidia
[conda] libblas                   3.9.0            14_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            14_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            14_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            14_linux64_mkl    conda-forge
[conda] mkl                       2022.0.1           h06a4308_117  
[conda] mkl-devel                 2022.0.1           h66538d2_117  
[conda] mkl-include               2022.0.1           h06a4308_117  
[conda] mypy                      0.942            py39hb9d737c_0    conda-forge
[conda] mypy_extensions           0.4.3            py39h06a4308_1  
[conda] numexpr                   2.8.0           mkl_py39h25e7801_2    conda-forge
[conda] numpy                     1.21.6           py39h18676bf_0    conda-forge
[conda] numpydoc                  1.3.1              pyhd8ed1ab_0    conda-forge
[conda] pytorch                   1.11.0          py3.9_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-forecasting       0.10.1             pyhd8ed1ab_0    conda-forge
[conda] pytorch-lightning         1.6.2              pyhd8ed1ab_0    conda-forge
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torch-tb-profiler         0.4.0                    pypi_0    pypi
[conda] torchaudio                0.11.0               py39_cu113    pytorch
[conda] torchdiffeq               0.2.2              pyhd8ed1ab_0    conda-forge
[conda] torchinfo                 1.6.5              pyhd8ed1ab_0    conda-forge
[conda] torchmetrics              0.8.1              pyhd8ed1ab_0    conda-forge
[conda] torchtext                 0.12.0                     py39    pytorch
[conda] torchvision               0.12.0               py39_cu113    pytorch
```