Fix API for OpenEquivariance.

Author: ameya98

src/e3tools/nn/__init__.py

@@ -4,7 +4,8 @@
     SeparableConv,
     SeparableConvBlock,
     FusedConv,
-    FusedDepthwiseConv,
+    FusedSeparableConv,
+    FusedSeparableConvBlock,
 )
 from ._linear import Linear
 from ._gate import Gate, Gated, GateWrapper
@@ -26,7 +27,8 @@
     "EquivariantMLP",
     "ExtractIrreps",
     "FusedConv",
-    "FusedDepthwiseConv",
+    "FusedSeparableConv",
+    "FusedSeparableConvBlock",
     "Gate",
     "GateWrapper",
     "Gated",

src/e3tools/nn/_conv.py

@@ -9,6 +9,7 @@
 from e3tools import scatter
 
 from ._gate import Gated
+from ._linear import Linear
 from ._interaction import LinearSelfInteraction
 from ._mlp import ScalarMLP
 from ._tensor_product import SeparableTensorProduct, DepthwiseTensorProduct
@@ -77,6 +78,8 @@ def __init__(
                 ```
             is used.
         """
+        if not openequivariance_available:
+            raise ImportError(f"OpenEquivariance could not be imported:\n{error_msg}")
 
         super().__init__()
 
@@ -101,20 +104,24 @@ def __init__(
 
         self.radial_nn = radial_nn(edge_attr_dim, self.tp.weight_numel)
 
-        if not openequivariance_available:
-            raise ImportError(f"OpenEquivariance could not be imported:\n{error_msg}")
+        if isinstance(self.tp, SeparableTensorProduct):
+            tp = self.tp.dtp
+            self.has_post_linear = True
+        else:
+            tp = self.tp
+            self.has_post_linear = False
 
-        # Remove path weight and path shape from instructions.
-        oeq_instructions = [instruction[:5] for instruction in self.tp.instructions]
+        # Remove path weight and path shape from instructions for OpenEquivariance.
+        oeq_instructions = [instruction[:5] for instruction in tp.instructions]
         oeq_tpp = oeq.TPProblem(
-            self.tp.irreps_in1,
-            self.tp.irreps_in2,
-            self.tp.irreps_out,
+            tp.irreps_in1,
+            tp.irreps_in2,
+            tp.irreps_out,
             oeq_instructions,
             shared_weights=False,
             internal_weights=False,
         )
-        self.fused_tp = oeq.TensorProductConv(
+        self.fused_tp_conv = oeq.TensorProductConv(
             oeq_tpp, torch_op=True, deterministic=False, use_opaque=False
         )
 
@@ -145,7 +152,10 @@ def forward(
 
         src, dst = edge_index
         radial_attr = self.radial_nn(edge_attr)
-        messages_agg = self.fused_tp(node_attr, edge_sh, radial_attr, dst, src)
+        messages_agg = self.fused_tp_conv(node_attr, edge_sh, radial_attr, dst, src)
+        if self.has_post_linear:
+            messages_agg = self.tp.lin(messages_agg)
+
         num_neighbors = scatter(
             torch.ones_like(src), src, dim=0, dim_size=N, reduce="sum"
         )
@@ -287,9 +297,9 @@ def __init__(self, *args, **kwargs):
         )
 
 
-class FusedDepthwiseConv(FusedConv):
+class FusedSeparableConv(FusedConv):
     """
-    Equivariant convolution layer using separable tensor product
+    Equivariant convolution layer using separable tensor product, with fused OpenEquivariance kernels.
 
     ref: https://arxiv.org/abs/1802.08219
     ref: https://arxiv.org/abs/2206.11990
@@ -299,7 +309,7 @@ def __init__(self, *args, **kwargs):
         super().__init__(
             *args,
             **kwargs,
-            tensor_product=DepthwiseTensorProduct,
+            tensor_product=SeparableTensorProduct,
         )
 
 
@@ -401,5 +411,22 @@ def __init__(self, *args, **kwargs):
         super().__init__(
             *args,
             **kwargs,
-            conv=SeparableConv,  # Explicitly set the convolution type to SeparableConv
+            conv=SeparableConv,
+        )
+
+
+class FusedSeparableConvBlock(ConvBlock):
+    """e3tools.nn.ConvBlock with FusedSeparableConv as the underlying convolution layer."""
+
+    def __init__(self, *args, **kwargs):
+        """
+        Initializes the SeparableConvBlock.
+
+        All arguments are passed directly to the parent ConvBlock,
+        with the 'conv' argument specifically set to SeparableConv.
+        """
+        super().__init__(
+            *args,
+            **kwargs,
+            conv=FusedSeparableConv,
         )

src/e3tools/nn/_tensor_product.py

@@ -48,6 +48,8 @@ def __init__(
             internal_weights=False,
             shared_weights=False,
         )
+
+        # For book-keeping.
         self.irreps_out = self.tp.irreps_out
         self.weight_numel = self.tp.weight_numel
         self.instructions = self.tp.instructions
@@ -85,7 +87,6 @@ def __init__(
         self.lin = Linear(self.dtp.irreps_out, self.irreps_out)
 
         # For book-keeping.
-        self.instructions = self.dtp.instructions
         self.weight_numel = self.dtp.weight_numel
 
     def forward(self, x, y, weight):

tests/test_fused_conv.py

@@ -5,25 +5,33 @@
 from torch import nn
 import e3nn
 
-from e3tools.nn import Conv, FusedConv, DepthwiseTensorProduct, ScalarMLP
+from e3tools.nn import Conv, FusedConv, DepthwiseTensorProduct, ScalarMLP, SeparableConv, FusedSeparableConv, SeparableTensorProduct
 from e3tools import radius_graph
 
-TENSOR_PRODUCTS = [
-    functools.partial(
-        e3nn.o3.FullyConnectedTensorProduct,
-        shared_weights=False,
-        internal_weights=False,
-    ),
-    DepthwiseTensorProduct,
-]
 
-
-@pytest.mark.parametrize("tensor_product", TENSOR_PRODUCTS)
+@pytest.mark.parametrize("tensor_product_type", [
+    "default",
+    "depthwise",
+    "separable",
+])
 @pytest.mark.parametrize("seed", [0, 1])
-def test_fused_conv(tensor_product, seed):
+def test_fused_conv(tensor_product_type: str, seed: int):
     if not torch.cuda.is_available():
         pytest.skip("CUDA is not available")
 
+    if tensor_product_type == "default":
+        tensor_product = functools.partial(
+            e3nn.o3.FullyConnectedTensorProduct,
+            shared_weights=False,
+            internal_weights=False,
+        )
+
+    elif tensor_product_type == "depthwise":
+        tensor_product = DepthwiseTensorProduct
+
+    elif tensor_product_type == "separable":
+        tensor_product = SeparableTensorProduct
+
     torch.manual_seed(seed)
     torch.set_default_device("cuda")
 
@@ -33,34 +41,24 @@ def test_fused_conv(tensor_product, seed):
     irreps_in = e3nn.o3.Irreps("10x0e + 4x1o + 1x2e")
     irreps_sh = irreps_in.spherical_harmonics(2)
 
-    tp = tensor_product(irreps_in, irreps_sh, irreps_in)
-    common_radial_nn = ScalarMLP(
-        in_features=edge_attr_dim,
-        out_features=tp.weight_numel,
-        hidden_features=[edge_attr_dim],
-        activation_layer=nn.SiLU,
-    )
-
-    def radial_nn(edge_attr_dim: int, num_elements: int) -> nn.Module:
-        return common_radial_nn
-
     layer = Conv(
         irreps_in=irreps_in,
         irreps_out=irreps_in,
         irreps_sh=irreps_sh,
-        radial_nn=radial_nn,
         edge_attr_dim=edge_attr_dim,
         tensor_product=tensor_product,
     )
     fused_layer = FusedConv(
         irreps_in=irreps_in,
         irreps_out=irreps_in,
         irreps_sh=irreps_sh,
-        radial_nn=radial_nn,
         edge_attr_dim=edge_attr_dim,
         tensor_product=tensor_product,
     )
 
+    # Copy weights.
+    fused_layer.load_state_dict(layer.state_dict())
+
     pos = torch.randn(N, 3)
     node_attr = layer.irreps_in.randn(N, -1)