fixed various semantic errors

apax/nn/torch/layers/__init__.py

@@ -1,3 +1,3 @@
-from apax.nn.torch.layers import activation, descriptor, ntk_linear, scaling
+from apax.nn.torch.layers import descriptor, ntk_linear, scaling
 
-__all__ = ["descriptor", "activation", "ntk_linear", "scaling"]
+__all__ = ["descriptor", "ntk_linear", "scaling"]

apax/nn/torch/layers/activation.py

@@ -0,0 +1,13 @@
+import torch.nn as nn
+from apax.nn.impl.activation import swish
+
+
+class SwishT(nn.Module):
+    def __init__(
+        self
+    ) -> None:
+        super().__init__()
+
+    def forward(self, x):
+        h = swish(x)
+        return h

apax/nn/torch/layers/descriptor/basis.py

@@ -62,8 +62,9 @@ def __init__(
         self.embeddings = None
         if self.emb_init is not None:
             self.embeddings = nn.Parameter()
+            self.n_radial = n_radial
         else:
-            self._n_radial = self.basis_fn.n_basis
+            self.n_radial = self.basis_fn.n_basis
 
     def forward(self, dr, Z_i, Z_j):
         dr = dr.type(self.dtype)

apax/nn/torch/layers/descriptor/gaussian_moment_descriptor.py

@@ -29,7 +29,7 @@ def __init__(
         self.dtype = dtype
 
         self.r_max = self.radial_fn.r_max
-        self.n_radial = self.radial_fn._n_radial
+        self.n_radial = self.radial_fn.n_radial
 
         self.distance = torch.vmap(distance, 0, 0)
 

apax/nn/torch/layers/readout.py

@@ -3,23 +3,23 @@
 import torch.nn.functional as F
 from typing import Any, Callable, List
 
-from apax.nn.impl.activation import swish
+from apax.nn.torch.layers.activation import SwishT
 from apax.nn.torch.layers.ntk_linear import NTKLinear
 
 
 class AtomisticReadoutT(nn.Module):
     def __init__(
-        self, units: List[int] = [512, 512], activation_fn: Callable = swish
+        self, units: List[int] = [512, 512], activation_fn: Callable = SwishT
     ) -> None:
         super().__init__()
 
-        units = [u for u in self.units] + [1]
+        units = [u for u in units] + [1]
         dense = []
         for ii, n_hidden in enumerate(units):
             dense.append(NTKLinear(n_hidden))
             if ii < len(units) - 1:
-                dense.append(activation_fn)
-        self.sequential = nn.Sequential(dense)
+                dense.append(activation_fn())
+        self.sequential = nn.Sequential(*dense)
 
     def forward(self, x):
         h = self.sequential(x)

apax/nn/torch/layers/scaling.py

@@ -15,8 +15,8 @@ def __init__(
         super().__init__()
         self.n_species = n_species
 
-        self.scale_param = nn.Parameter(scale)
-        self.shift_param = nn.Parameter(shift)
+        self.scale_param = nn.Parameter(torch.tensor(scale))
+        self.shift_param = nn.Parameter(torch.tensor(shift))
         self.dtype = dtype
 
     def forward(self, x: torch.Tensor, Z: torch.Tensor) -> torch.Tensor:

apax/ops.py

@@ -58,10 +58,10 @@ def sum(x, axis, keepdims=False):
         return torch.sum(x, dim=axis, keepdim=keepdims)
 
 
-def concatenate(x, axis):
-    if isinstance(x, jax.Array):
+def concatenate(x: list, axis):
+    if isinstance(x[0], jax.Array):
         return jax.numpy.concatenate(x, axis=axis)
-    elif isinstance(x, torch.Tensor):
+    elif isinstance(x[0], torch.Tensor):
         return torch.concatenate(x, dim=axis)