Circular imports when importing einops and torch._dynamo about einops HOT 6 OPEN

Hi @befelix , thank you for reporting, there is a circular import problem that we need to address.

step-by-step what happens in python -W "error" -c 'import einops; import torch._dynamo'

1. einops is imported
2. torch._dynamo is imported, it detects that einops is imported, and calls https://github.com/pytorch/pytorch/blob/88a71594933b2464d9d8b6b3533c5a945a4ac2ff/torch/_dynamo/decorators.py#L322
3. ... which imports from einops._torch_specific ...
4. einops._torch_specific tries to register ops by using allow_in_graph, that is needs to import first from torch._dynamo. That's where we run in circular imports

cc @wconstab @jansel

from einops.

Seems this torch PR was trying to address the problem, but that didn't quite work:
https://github.com/pytorch/pytorch/pull/111835/files

Fix wouldn't work it seems: einops needs allow_in_graph being importable, meaning that hook can only be called after _dynamo was imported. Currently hook is called during the import.

from einops.

Should we delete the special einops handling from torch._dynamo?

from einops.

it is preferrable to defer that logic to 'on first compile' (or when torch._dynamo is fully imported)

Otherwise, we better just delete that handling code - it wouldn't work as intended

from einops.

Can you check if this fixes it: pytorch/pytorch#124084

from einops.

Thank you @jansel for merging that. I can confirm nightly torch==2.4.0.dev20240420+cpu does not have this issue.

Before fix comes live users need to use manual registration of functions:
https://github.com/arogozhnikov/einops/wiki/Using-torch.compile-with-einops

from einops.