Fix TorchScript warning "doesn't support instance-level annotations" (e3nn#437)

Author: Linux-cpp-lisp

e3nn/o3/__init__.py

@@ -103,7 +103,8 @@
     "FullyConnectedTensorProduct",
     "ElementwiseTensorProduct",
     "FullTensorProduct",
-    "FullTensorProductv2" "TensorSquare",
+    "FullTensorProductv2",
+    "TensorSquare",
     "SphericalHarmonics",
     "spherical_harmonics",
     "SphericalHarmonicsAlphaBeta",

e3nn/o3/experimental/__init__.py

@@ -1 +1,3 @@
 from ._full_tp import FullTensorProduct as FullTensorProductv2
+
+__all__ = [FullTensorProductv2]

e3nn/o3/experimental/_full_tp.py

@@ -1,3 +1,5 @@
+# flake8: noqa
+
 from e3nn.util.datatypes import Path, Chunk
 from e3nn import o3
 

e3nn/util/codegen/_mixin.py

@@ -1,9 +1,8 @@
 import io
 from typing import Dict
 
-import e3nn
+import e3nn.util.jit
 import torch
-from opt_einsum_fx import jitable
 from torch import fx
 
 
@@ -40,7 +39,9 @@ def _codegen_register(
             assert isinstance(graphmod, fx.GraphModule)
 
             if opt_defaults["jit_script_fx"]:
-                scriptmod = torch.jit.script(jitable(graphmod))
+                # With recurse=False, this more or less is equivalent to
+                # torch.jit.script(jitable(graphmod))
+                scriptmod = e3nn.util.jit.compile(graphmod, recurse=False)
                 assert isinstance(scriptmod, torch.jit.ScriptModule)
             else:
                 scriptmod = graphmod
@@ -73,7 +74,7 @@ def __getstate__(self):
                 # Get the module
                 smod = getattr(self, fname)
                 if isinstance(smod, fx.GraphModule):
-                    smod = torch.jit.script(jitable(smod))
+                    smod = e3nn.util.jit.compile(smod, recurse=False)
                 assert isinstance(smod, torch.jit.ScriptModule)
                 # Save the compiled code as TorchScript IR
                 buffer = io.BytesIO()

e3nn/util/jit.py

@@ -1,6 +1,7 @@
 import copy
 import inspect
 import warnings
+import re
 from typing import Optional
 
 import torch
@@ -59,6 +60,7 @@ def compile(
     script_options: dict = None,
     trace_options: dict = None,
     in_place: bool = True,
+    recurse: bool = True,
 ):
     """Recursively compile a module and all submodules according to their decorators.
 
@@ -76,6 +78,10 @@ def compile(
             Extra kwargs for ``torch.jit.script``.
         trace_options : dict, default = {}
             Extra kwargs for ``torch.jit.trace``.
+        in_place : bool, default True
+            Whether to insert the recursively compiled submodules in-place, or do a deepcopy first.
+        recurse : bool, default True
+            Whether to recurse through the module's children before passing the parent to TorchScript
 
     Returns
     -------
@@ -92,25 +98,57 @@ def compile(
         mod = copy.deepcopy(mod)
     # TODO: debug logging
     assert n_trace_checks >= 1
-    # == recurse to children ==
-    # This allows us to trace compile submodules of modules we are going to script
-    for submod_name, submod in mod.named_children():
-        setattr(
-            mod,
-            submod_name,
-            compile(
-                submod,
-                n_trace_checks=n_trace_checks,
-                script_options=script_options,
-                trace_options=trace_options,
-                in_place=True,  # since we deepcopied the module above, we can do inplace
-            ),
-        )
+
+    if recurse:
+        # == recurse to children ==
+        # This allows us to trace compile submodules of modules we are going to script
+        for submod_name, submod in mod.named_children():
+            setattr(
+                mod,
+                submod_name,
+                compile(
+                    submod,
+                    n_trace_checks=n_trace_checks,
+                    script_options=script_options,
+                    trace_options=trace_options,
+                    in_place=True,  # since we deepcopied the module above, we can do inplace
+                    recurse=recurse,  # always true in this branch
+                ),
+            )
+
     # == Compile this module now ==
     if mode == "script":
         if isinstance(mod, fx.GraphModule):
             mod = jitable(mod)
-        mod = torch.jit.script(mod, **script_options)
+            # In recent PyTorch versions (probably >1.12, definitely >=2.0), PyTorch's implementation of fx.GraphModule
+            # causes a warning to be raised when fx.GraphModules are compiled to TorchScript with `torch.jit.script`:
+            #
+            #   torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level
+            #   annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the
+            #   class body, or 2) wrap the type in `torch.jit.Attribute`.
+            #
+            # Using the debugger traces this back to the following line in PyTorch:
+            # https://github.com/pytorch/pytorch/blob/v2.3.1/torch/fx/graph_module.py#L446
+            # Because the metadata stored by GraphModule is not relevant to the compiled TorchScript module
+            # we need, it should be safe to ignore this warning. The below code suppresses this warning as
+            # narrowly as possible to ensure it can still be raised from user code.
+            # See also: https://github.com/pytorch/pytorch/issues/89064
+            with warnings.catch_warnings():
+                warnings.filterwarnings(
+                    "ignore",
+                    # warnings treats this argument as a regex, but we want to match a string literal exactly, so escape it:
+                    message=re.escape(
+                        "The TorchScript type system doesn't support instance-level annotations on empty non-base types "
+                        "in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type "
+                        "in `torch.jit.Attribute`."
+                    ),
+                    # Being specific is good form, even though matching the message should be enough:
+                    category=UserWarning,
+                    module="torch",
+                )
+                mod = torch.jit.script(mod, **script_options)
+        else:
+            mod = torch.jit.script(mod, **script_options)
     elif mode == "trace":
         # These are always modules, so we're always using trace_module
         # We need tracing inputs:

examples/conftest.py

@@ -12,7 +12,7 @@ def test_jit(float_tolerance) -> None:
 
     a = torch.randn(5, 4)
     b = torch.randn(5, 4)
-    return (sh(a, b) - jited(a, b)).abs().max() < float_tolerance
+    assert (sh(a, b) - jited(a, b)).abs().max() < float_tolerance
 
 
 def test_sh_equivariance1(float_tolerance) -> None:

tests/o3/angular_spherical_harmonics_test.py

@@ -1,28 +1,32 @@
-import torch
-from torch._inductor.utils import print_performance
+# flake8: noqa
 
-# Borrowed from https://github.com/pytorch-labs/gpt-fast/blob/db7b273ab86b75358bd3b014f1f022a19aba4797/generate.py#L16-L18
-torch.set_float32_matmul_precision("high")
-import torch._dynamo.config
-import torch._inductor.config
 
-torch._inductor.config.coordinate_descent_tuning = True
-torch._inductor.config.triton.unique_kernel_names = True
+def main():
+    import torch
+    from torch._inductor.utils import print_performance
 
-device = "cuda"
-compile_mode = "max-autotune"  # Bringing out all of the tricks that Torch 2.0 has but "reduce-overhead" should work as well
+    # Borrowed from https://github.com/pytorch-labs/gpt-fast/blob/db7b273ab86b75358bd3b014f1f022a19aba4797/generate.py#L16-L18
+    torch.set_float32_matmul_precision("high")
+    import torch._dynamo.config
+    import torch._inductor.config
 
-from e3nn import o3, util
-import numpy as np
-from torch import nn
-import time
+    torch._inductor.config.coordinate_descent_tuning = True
+    torch._inductor.config.triton.unique_kernel_names = True
 
-LMAX = 8
-CHANNEL = 128
-BATCH = 100
+    device = "cuda"
+    compile_mode = (
+        "max-autotune"  # Bringing out all of the tricks that Torch 2.0 has but "reduce-overhead" should work as well
+    )
 
+    from e3nn import o3, util
+    import numpy as np
+    from torch import nn
+    import time
+
+    LMAX = 8
+    CHANNEL = 128
+    BATCH = 100
 
-def main():
     for lmax in range(1, LMAX + 1):
         irreps = o3.Irreps.spherical_harmonics(lmax)
         irreps_x = (CHANNEL * irreps).regroup()