Make openeq import only with FusedConv. Add Repeat().

Author: ameya98

src/e3tools/nn/__init__.py

@@ -17,6 +17,7 @@
 from ._transformer import Attention, MultiheadAttention, TransformerBlock
 from ._extract_irreps import ExtractIrreps
 from ._scaling import ScaleIrreps
+from ._repeat import Repeat
 
 __all__ = [
     "Attention",
@@ -37,6 +38,7 @@
     "LinearSelfInteraction",
     "MulToAxis",
     "MultiheadAttention",
+    "Repeat",
     "ScalarMLP",
     "ScaleIrreps",
     "SeparableConv",

src/e3tools/nn/_conv.py

@@ -13,14 +13,6 @@
 from ._mlp import ScalarMLP
 from ._tensor_product import SeparableTensorProduct
 
-try:
-    import openequivariance as oeq
-
-    openequivariance_available = True
-except ImportError as e:
-    error_msg = str(e)
-    openequivariance_available = False
-
 
 class FusedConv(nn.Module):
     """
@@ -77,8 +69,10 @@ def __init__(
                 ```
             is used.
         """
-        if not openequivariance_available:
-            raise ImportError(f"OpenEquivariance could not be imported:\n{error_msg}")
+        try:
+            import openequivariance as oeq
+        except ImportError as e:
+            raise ImportError(f"OpenEquivariance could not be imported: {e}")
 
         super().__init__()
 

src/e3tools/nn/_repeat.py

@@ -0,0 +1,24 @@
+import torch
+from torch import nn
+import e3nn.o3
+
+
+
+class Repeat(nn.Module):
+    """Repeat the irreps along the last axis."""
+
+    def __init__(self, irreps_in: e3nn.o3.Irreps, repeats: int):
+        super().__init__()
+        self.irreps_in = e3nn.o3.Irreps(irreps_in)
+        self.repeats = repeats
+        self.irreps_out = e3nn.o3.Irreps(
+            [(mul * repeats, ir) for mul, ir in irreps_in]
+        )
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Repeat the features along the last axis.
+
+        If `x` has shape `[..., irreps_in.dim]`, the output will have shape
+        `[..., repeats * irreps_in.dim]`.
+        """
+        return x.repeat_interleave(self.repeats, dim=-1)