[Dynamo] only import einops if version is lower than 0.7.0 (#142847)

mlazos · pytorchmergebot · commit 357e261b1ede · 2024-12-12T06:38:22.000Z
Fixes internal xref (https://fb.workplace.com/groups/257735836456307/posts/804793021750583/?comment_id=805229281706957&reply_comment_id=805232695039949) Pull Request resolved: #142847 Approved by: https://github.com/zou3519
diff --git a/torch/_dynamo/decorators.py b/torch/_dynamo/decorators.py
@@ -2,9 +2,12 @@
 # ruff: noqa: TCH004
 import functools
 import inspect
+import sys
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, Type, TYPE_CHECKING, TypeVar
 
+from packaging import version
+
 import torch
 from torch.utils._contextlib import _DecoratorContextManager
 from torch.utils._python_dispatch import is_traceable_wrapper_subclass
@@ -608,27 +611,33 @@ def mark_static_address(t, guard=True):
 # Note: this carefully avoids eagerly import einops.
 # TODO: we should delete this whole _allow_in_graph_einops logic by approximately 2024 Q2
 def _allow_in_graph_einops():
-    import einops
-
-    try:
-        # requires einops > 0.6.1, torch >= 2.0
-        from einops._torch_specific import (  # type: ignore[attr-defined]  # noqa: F401
-            _ops_were_registered_in_torchdynamo,
-        )
+    mod = sys.modules.get("einops")
+    if mod is None:
+        return
+    else:
+        # version > 0.7.0 does allow_in_graph out of tree
+        if version.Version(mod.__version__) < version.Version("0.7.0"):
+            import einops
 
-        # einops > 0.6.1 will call the op registration logic as it is imported.
-    except ImportError:
-        # einops <= 0.6.1
-        allow_in_graph(einops.rearrange)
-        allow_in_graph(einops.reduce)
-        if hasattr(einops, "repeat"):
-            allow_in_graph(einops.repeat)  # available since einops 0.2.0
-        if hasattr(einops, "einsum"):
-            allow_in_graph(einops.einsum)  # available since einops 0.5.0
-        if hasattr(einops, "pack"):
-            allow_in_graph(einops.pack)  # available since einops 0.6.0
-        if hasattr(einops, "unpack"):
-            allow_in_graph(einops.unpack)  # available since einops 0.6.0
+            try:
+                # requires einops > 0.6.1, torch >= 2.0
+                from einops._torch_specific import (  # type: ignore[attr-defined]  # noqa: F401
+                    _ops_were_registered_in_torchdynamo,
+                )
+
+                # einops > 0.6.1 will call the op registration logic as it is imported.
+            except ImportError:
+                # einops <= 0.6.1
+                allow_in_graph(einops.rearrange)
+                allow_in_graph(einops.reduce)
+                if hasattr(einops, "repeat"):
+                    allow_in_graph(einops.repeat)  # available since einops 0.2.0
+                if hasattr(einops, "einsum"):
+                    allow_in_graph(einops.einsum)  # available since einops 0.5.0
+                if hasattr(einops, "pack"):
+                    allow_in_graph(einops.pack)  # available since einops 0.6.0
+                if hasattr(einops, "unpack"):
+                    allow_in_graph(einops.unpack)  # available since einops 0.6.0
 
 
 trace_rules.add_module_init_func("einops", _allow_in_graph_einops)
