Resolve #126

Co-authored-by: Félix <[email protected]>

nequip/nn/_convnetlayer.py

@@ -10,18 +10,10 @@
     GraphModuleMixin,
     InteractionBlock,
 )
-from nequip.nn.nonlinearities import ShiftedSoftPlus
+from nequip.nn.nonlinearities import get_nonlinearity
 from nequip.utils.tp_utils import tp_path_exists
 
 
-acts = {
-    "abs": torch.abs,
-    "tanh": torch.tanh,
-    "ssp": ShiftedSoftPlus,
-    "silu": torch.nn.functional.silu,
-}
-
-
 class ConvNetLayer(GraphModuleMixin, torch.nn.Module):
     """
     Args:
@@ -96,15 +88,16 @@ def __init__(
             )
             irreps_gates = o3.Irreps([(mul, ir) for mul, _ in irreps_gated])
 
-            # TO DO, it's not that safe to directly use the
-            # dictionary
             equivariant_nonlin = Gate(
                 irreps_scalars=irreps_scalars,
                 act_scalars=[
-                    acts[nonlinearity_scalars[ir.p]] for _, ir in irreps_scalars
+                    get_nonlinearity(nonlinearity_scalars[ir.p])
+                    for _, ir in irreps_scalars
                 ],
                 irreps_gates=irreps_gates,
-                act_gates=[acts[nonlinearity_gates[ir.p]] for _, ir in irreps_gates],
+                act_gates=[
+                    get_nonlinearity(nonlinearity_gates[ir.p]) for _, ir in irreps_gates
+                ],
                 irreps_gated=irreps_gated,
             )
 
@@ -116,7 +109,7 @@ def __init__(
             equivariant_nonlin = NormActivation(
                 irreps_in=conv_irreps_out,
                 # norm is an even scalar, so use nonlinearity_scalars[1]
-                scalar_nonlinearity=acts[nonlinearity_scalars[1]],
+                scalar_nonlinearity=get_nonlinearity(nonlinearity_scalars[1]),
                 normalize=True,
                 epsilon=1e-8,
                 bias=False,

nequip/nn/_interaction_block.py

@@ -10,7 +10,7 @@
 from e3nn.o3 import TensorProduct, Linear, FullyConnectedTensorProduct
 
 from nequip.data import AtomicDataDict
-from nequip.nn.nonlinearities import ShiftedSoftPlus
+from nequip.nn.nonlinearities import get_nonlinearity
 from ._graph_mixin import GraphModuleMixin
 
 
@@ -115,10 +115,7 @@ def __init__(
             [self.irreps_in[AtomicDataDict.EDGE_EMBEDDING_KEY].num_irreps]
             + invariant_layers * [invariant_neurons]
             + [tp.weight_numel],
-            {
-                "ssp": ShiftedSoftPlus,
-                "silu": torch.nn.functional.silu,
-            }[nonlinearity_scalars["e"]],
+            get_nonlinearity(nonlinearity_scalars["e"]),
         )
 
         self.tp = tp

nequip/nn/nonlinearities.py

@@ -1,8 +1,31 @@
-import torch
+from typing import Callable
 
 import math
 
+import torch
+
+from e3nn.util.jit import compile_mode
+
 
 @torch.jit.script
-def ShiftedSoftPlus(x):
+def shifted_soft_plus(x: torch.Tensor) -> torch.Tensor:
     return torch.nn.functional.softplus(x) - math.log(2.0)
+
+
+@compile_mode("script")
+class ShiftedSoftPlus(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return shifted_soft_plus(x)
+
+
+def get_nonlinearity(name: str) -> Callable:
+    if name == "abs":
+        return torch.abs
+    elif name == "tanh":
+        return torch.nn.Tanh()
+    elif name == "ssp":
+        return ShiftedSoftPlus()
+    elif name == "silu":
+        return torch.nn.SiLU()
+    else:
+        raise KeyError(f"No such nonlinearity `{name}`")