Skip allow_in_graph if torch >= 2.8 (#394)

Starting from PyTorch 2.8, einops no longer needs to rely on `allow_in_graph` for its own ops. PyTorch follows the same approach internally, and I'd like to remove this usage from the PyTorch codebase as well. Dynamo is now mature enough to trace through the einops codebase, and we have tests in place to catch any regressions early.

Author: guilhermeleobas

einops/_torch_specific.py

@@ -104,6 +104,11 @@ def allow_ops_in_compiled_graph():
     if hasattr(torch, "__version__") and torch.__version__[0] < "2":
         # torch._dynamo and torch.compile appear in pytorch 2.0
         return
+
+    if hasattr(torch, "__version__") and torch.__version__ >= "2.8":
+        # einops don't need to use allow_in graph for torch 2.8 and above
+        return
+
     try:
         from torch._dynamo import allow_in_graph
     except ImportError: