Merge pull request #8 from Ipsedo/develop

change conv weights init + einsum for linear
Ipsedo · Jul 15, 2024 · 6d962ab · 6d962ab
2 parents 8645ee3 + 2360521
commit 6d962ab
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/kan/networks/conv.py b/kan/networks/conv.py
@@ -42,9 +42,8 @@ def __init__(
             )
         )
 
-        xavier_normal_(self._w_b, 1.0 / kernel_size**2)
-        # xavier_normal_(self._w_s, 1e-3)
-        normal_(self._c, 0, 1e-1 / kernel_size**2)
+        xavier_normal_(self._w_b)
+        normal_(self._c, 0, 1e-1)
 
         self._in_channels = in_channels
         self._kernel_size = kernel_size

diff --git a/kan/networks/linear.py b/kan/networks/linear.py
@@ -34,7 +34,7 @@ def forward(self, x: th.Tensor) -> th.Tensor:
         assert len(x.size()) == 2
 
         return th.sum(
-            self._w_s * th.sum(self.__act_fun(x).unsqueeze(2) * self._c, dim=1)
-            + self._w_b * self.__res_act_fun(x).unsqueeze(1),
+            self._w_s * th.einsum("bai,aoi->boi", self.__act_fun(x), self._c)
+            + th.einsum("bi,oi->boi", self.__res_act_fun(x), self._w_b),
             dim=2,
         )