Refactor input checks in tensor product classes to include tracing condition

Author: mariogeiger

cuequivariance_torch/cuequivariance_torch/primitives/equivariant_tensor_product.py

@@ -251,7 +251,11 @@ def forward(
         """
         If ``indices`` is not None, the first input is indexed by ``indices``.
         """
-        if not torch.jit.is_scripting() and not torch.compiler.is_compiling():
+        if (
+            not torch.jit.is_scripting()
+            and not torch.jit.is_tracing()
+            and not torch.compiler.is_compiling()
+        ):
             if not isinstance(inputs, (list, tuple)):
                 raise ValueError(
                     "inputs should be a list of tensors followed by optional indices"

cuequivariance_torch/cuequivariance_torch/primitives/symmetric_tensor_product.py

@@ -320,7 +320,11 @@ def forward(
         i0 = i0.to(torch.int32)
         x0 = x0.reshape(x0.shape[0], x0.shape[1] // self.u, self.u)
         x1 = x1.reshape(x1.shape[0], x1.shape[1] // self.u, self.u)
-        if not torch.jit.is_scripting() and not torch.compiler.is_compiling():
+        if (
+            not torch.jit.is_scripting()
+            and not torch.jit.is_tracing()
+            and not torch.compiler.is_compiling()
+        ):
             logger.debug(
                 f"Calling SymmetricTensorContraction: {self.descriptors}, input shapes: {x0.shape}, {i0.shape}, {x1.shape}"
             )

cuequivariance_torch/cuequivariance_torch/primitives/tensor_product.py

@@ -116,7 +116,11 @@ def forward(self, inputs: List[torch.Tensor]):
                 It has a shape of (batch, last_operand_size), where
                 `last_operand_size` is the size of the last operand in the descriptor.
         """
-        if not torch.jit.is_scripting() and not torch.compiler.is_compiling():
+        if (
+            not torch.jit.is_scripting()
+            and not torch.jit.is_tracing()
+            and not torch.compiler.is_compiling()
+        ):
             if not isinstance(inputs, (list, tuple)):
                 raise ValueError("inputs should be a list of tensors")
             if len(inputs) != self.num_operands - 1:
@@ -374,7 +378,11 @@ def __init__(self, module: torch.nn.Module, descriptor: stp.SegmentedTensorProdu
         self.descriptor = descriptor
 
     def forward(self, args: List[torch.Tensor]):
-        if not torch.jit.is_scripting() and not torch.compiler.is_compiling():
+        if (
+            not torch.jit.is_scripting()
+            and not torch.jit.is_tracing()
+            and not torch.compiler.is_compiling()
+        ):
             for oid, arg in enumerate(args):
                 torch._assert(
                     arg.ndim == 2,
@@ -506,7 +514,11 @@ def __repr__(self) -> str:
     def forward(self, inputs: List[torch.Tensor]) -> torch.Tensor:
         x0, x1 = self._perm(inputs[0], inputs[1])
 
-        if not torch.jit.is_scripting() and not torch.compiler.is_compiling():
+        if (
+            not torch.jit.is_scripting()
+            and not torch.jit.is_tracing()
+            and not torch.compiler.is_compiling()
+        ):
             logger.debug(
                 f"Calling FusedTensorProductOp3: {self.descriptor}, input shapes: {x0.shape}, {x1.shape}"
             )
@@ -567,7 +579,11 @@ def __repr__(self) -> str:
     def forward(self, inputs: List[torch.Tensor]) -> torch.Tensor:
         x0, x1, x2 = self._perm(inputs[0], inputs[1], inputs[2])
 
-        if not torch.jit.is_scripting() and not torch.compiler.is_compiling():
+        if (
+            not torch.jit.is_scripting()
+            and not torch.jit.is_tracing()
+            and not torch.compiler.is_compiling()
+        ):
             logger.debug(
                 f"Calling FusedTensorProductOp4: {self.descriptor}, input shapes: {x0.shape}, {x1.shape}, {x2.shape}"
             )
@@ -626,7 +642,11 @@ def __repr__(self):
     def forward(self, inputs: List[torch.Tensor]) -> torch.Tensor:
         x0, x1 = inputs
 
-        if not torch.jit.is_scripting() and not torch.compiler.is_compiling():
+        if (
+            not torch.jit.is_scripting()
+            and not torch.jit.is_tracing()
+            and not torch.compiler.is_compiling()
+        ):
             logger.debug(
                 f"Calling TensorProductUniform3x1d: {self.descriptor}, input shapes: {x0.shape}, {x1.shape}"
             )
@@ -647,7 +667,11 @@ def __repr__(self):
     def forward(self, inputs: List[torch.Tensor]):
         x0, x1, x2 = inputs
 
-        if not torch.jit.is_scripting() and not torch.compiler.is_compiling():
+        if (
+            not torch.jit.is_scripting()
+            and not torch.jit.is_tracing()
+            and not torch.compiler.is_compiling()
+        ):
             logger.debug(
                 f"Calling TensorProductUniform4x1d: {self.descriptor}, input shapes: {x0.shape}, {x1.shape}, {x2.shape}"
             )