einops does not support torch.jit.script ?

[ahatamiz]

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.

[arogozhnikov]

Edit:

ignore my comment below. I didn't see you use module. Does it work for newer torch (torch.jit.script changed a lot since 1.6)?

---

Original comment:

torch.jit.script works with einops classes, not functions. Use class counterparts:

from einops.layers.torch import Rearrange, Reduce, Einmix

unfortunately, torch.jit.script-ing has multiple limitations that prevent scripting of functions.

There is a hacky workaround if you for some reason want to use functions: torch.jit.trace parts that contain einops functions, then torch.jit.script result. That said, I recommend using modules instead

[ahatamiz]

Hi @arogozhnikov

Yes. As seen in the snippet, a PyTorch module is used, and torch.jit.script is not supported.

Based on this conversation, I created a new environment with the latest PyTorch release (1.10.0). But it is still not supported as shown below.

I believe this is a highly important feature. As much as I like to include einops in my work, not being able to support torch.jit.script may force developers to reconsider this choice. As a result, it is appreciated if the issue can be re-opened and looked at.

Thanks again

[arogozhnikov]

@ahatamiz that's all strange, jit for modules is tested in CI.

I just installed latest pytorch and used your example in the first post, and it works.

Maybe you can specify your env then?

Here is what I had after a fresh install (ubuntu20.04):

!python --version
Python 3.10.4

import torch
torch.__version__

'1.12.1+cpu'

import einops
einops.__version__

'0.4.1'

[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)

[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.