rxnn 0.1.12__tar.gz → 0.1.13__tar.gz

{rxnn-0.1.12 → rxnn-0.1.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.1.12
+Version: 0.1.13
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning

{rxnn-0.1.12 → rxnn-0.1.13}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "rxnn"
-version = "0.1.12"
+version = "0.1.13"
 description = "RxNN: Reactive Neural Networks Platform"
 
 license = "Apache-2.0"

{rxnn-0.1.12 → rxnn-0.1.13}/src/rxnn/transformers/layers.py

@@ -59,7 +59,7 @@ class ReactiveTransformerLayer(nn.Module):
         for param in self.memory_cross_attention.parameters():
             param.requires_grad_(is_trainable)
 
-    def moe_router_loss_(self):
+    def moe_router_loss(self):
         return self.ff.router_loss() if self.use_moe else None
 
     def forward(self, x: torch.Tensor, stm: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
@@ -135,7 +135,7 @@ class ClassicTransformerLayer(nn.Module):
         self.use_post_norm = use_post_norm
         self.use_moe = use_moe
 
-    def moe_router_loss_(self):
+    def moe_router_loss(self):
         return self.ff.router_loss() if self.use_moe else torch.tensor(0.0)
 
     def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:

{rxnn-0.1.12 → rxnn-0.1.13}/src/rxnn/transformers/models.py

@@ -37,7 +37,7 @@ class ReactiveTransformerBase(nn.Module):
         for i in range(self.num_own_layers):
             self.layers[i].trainable_cross_attention_(is_trainable)
 
-    def moe_router_loss_(self):
+    def moe_router_loss(self):
         return torch.stack([self.layers[i].moe_router_loss() for i in range(self.num_own_layers) if self.layers[i].use_moe] + [
             self.shared_layers[i].moe_router_loss() for i in range(self.num_shared_layers) if self.shared_layers[i].use_moe]).mean()
 
@@ -123,7 +123,7 @@ class ClassicTransformerBase(nn.Module):
         self.layers = layers
         self.num_layers = len(layers) if layers else 0
 
-    def moe_router_loss_(self):
+    def moe_router_loss(self):
         return torch.stack([self.layers[i].moe_router_loss() for i in range(self.num_layers) if self.layers[i].use_moe]).mean()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor: