Open befelix opened 4 months ago
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'
from einops._torch_specific ...
allow_in_graph
, that is needs to import first from torch._dynamo. That's where we run in circular importscc @wconstab @jansel
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.
Should we delete the special einops handling from torch._dynamo
?
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
Can you check if this fixes it: https://github.com/pytorch/pytorch/pull/124084
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
this was fixed in pytorch 2.4, I'll leave the issue open until 2.5 comes out.
Describe the bug Importing
einops
beforetorch._dynamo
currently leads to warnings. I'm not sure if this needs a fix on thepytorch
oreinops
side. This is annoying for CI pipelines, where warnings are typically treated as errors. Note that, with sorted imports,einops
will typically import before anytorch
namespaces.Reproduction steps Steps to reproduce the behavior:
Expected behavior
No output, import order does not matter.
Your platform
Ubuntu 22.04,
torch==2.2.2
,einops==0.7.0