Identity skip + args for edge gated block

mace/blocks.py

@@ -90,6 +90,10 @@ def __init__(
         hidden_irreps: o3.Irreps,
         avg_num_neighbors: float,
         rbf_hidden_channels: int = 64,
+        edge_gates_irreps: o3.Irreps = o3.Irreps("0e"),
+        num_gates: int = 4,
+        multi_conv: bool = False,
+        exponential: bool = True,
     ) -> None:
         super().__init__()
         self.node_attrs_irreps = node_attrs_irreps
@@ -100,6 +104,10 @@ def __init__(
         self.hidden_irreps = hidden_irreps
         self.avg_num_neighbors = avg_num_neighbors
         self.rbf_hidden_channels = rbf_hidden_channels
+        self.edge_gates_irreps = edge_gates_irreps
+        self.num_gates = num_gates
+        self.multi_conv = multi_conv
+        self.exponential = exponential
 
         self._setup()
 

mace/irreps_tools.py

@@ -84,3 +84,17 @@ def forward(self, tensor: torch.Tensor) -> torch.Tensor:
             field = field.reshape(batch, mul, d)
             out.append(field)
         return torch.cat(out, dim=-1)
+
+
+@compile_mode("script")
+class lifted_skip(torch.nn.Module):
+    def __init__(self, irreps_in: o3.Irreps, irreps_out: o3.Irreps) -> None:
+        super().__init__()
+        self.irreps_in = o3.Irreps(irreps_in)
+        self.irreps_out = o3.Irreps(irreps_out)
+
+    def forward(self, tensor: torch.Tensor) -> torch.Tensor:
+        batch, _ = tensor.shape
+        template = torch.zeros(batch, self.irreps_out.dim, device=tensor.device)
+        template[:, 0 : self.irreps_in.dim] = tensor
+        return template