einops does not support torch.jit.script ?

Describe the bug Thanks for this great work. In recent updates, it is mentioned that einops supports torch.jit.script for PyTorch layers. I and @yiheng-wang-nv have been looking into this to support TorchScript for a number of models. However, we are not able to utilize this functionality for simple operations such as Rearrange.

Reproduction steps The following snippet should illustrate a concise way of reproducing this issue:

import torch
import torch.nn as nn
from einops.layers.torch import Rearrange

class SimpleRearrange(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Sequential(Rearrange('b c h w -> b h w c'))
    
    def forward(self, x):
        result = self.layer(x)
        return result
    
net = SimpleRearrange()
net.eval()
with torch.no_grad():
    torch.jit.script(net)

Expected behavior This is the expected output:

Your platform einops version: 0.3.2 Python version: 3.8.12 PyTorch version: 1.6.0 CUDA version: 10.2

Based on this, I believe einops does not support torch.jit.script, unless we are missing something. Appreciate your inputs here.

[alexfanqi]

@arogozhnikov Thanks for the kind tips. My mind was not in a proper state yesterday. The bug is indeed fixed in another issue. It is now so embarrassing to recall my moment yesterday.

[arogozhnikov]

Also regarding this deleted comment:

It seems I am able to work around the issue by decorating TransformRecipe with torch.jit.script. But I have no idea how to incorporate this into einops.

@torch.jit.script
class TransformRecipe:

If that is true (again, strange, as torch should do that for you), you should be able to use hot-patching:

from einops import einops as _einops
_einops.TransformRecipe = torch.jit.script(_einops.TransformRecipe)