[BUG] Updated: einops does not support torch.jit.script

[ahatamiz]

Updated
The original issue was closed, despite the fact that provided snippet still uses a module, as opposed to a function and tested with newer and older PyTorch versions. Opening the issue here in hope of restarting the conversation and finding possible solutions.

Describe the bug
The einops package does not support torch.jit.script. The following snippet tries a simple Rearrange operation, using a module to test this. It is observed that the operation is not supported.

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.10.0
CUDA version: 10.2

[arogozhnikov]

Hey, thanks for bringing up

This is strange, but I can't reproduce this on mac, and I can't reproduce this on linux.
There is also test in CI

einops/tests/test_layers.py

Line 227 in 85dcd35

model1 = create_torch_model(use_reduce=True, add_scripted_layer=script_layer)

I've tested 0.3.2 and my local master

I think following things should be checked:

1. einops.__file__ and einops.einops.__file__ point to the right location
2. verify that simple command is scriptable, like

torch.jit.script(Rearrange('b c h w -> b h w c'))

3. I have observed previously strange behavior when torch.jit.script for some reason fails on the first run of scripting, but after catching exception and running the same command again, it succeeds. Please check if this is the case, like

try:
    torch.jit.script(net) # fails
except:
    torch.jit.script(net) # succeeds, while nothing has changed

[arogozhnikov]

I mean - your snippet ran successfully for me in all cases

[ahatamiz]

Hi @arogozhnikov

Thanks for the reply. Again, I have tested the code snippet and getting the same error consistently. It would be great if this can be tested by another person.

[arogozhnikov]

added one more point to check to my previous comment ^

[ahatamiz]

Hi @arogozhnikov

It seems that TorchScript support is made available in this commit which is not included in 0.3.2.

Considering this, would it be possible to release a new version and include it ?

Thanks again
Ali

[arogozhnikov]

Hi Ali, thanks for tracking this down. I've published a new version to pypi, it contains this commit, please try it