rxnn 0.1.60__tar.gz → 0.1.61__tar.gz

{rxnn-0.1.60 → rxnn-0.1.61}/PKG-INFO

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

{rxnn-0.1.60 → rxnn-0.1.61}/pyproject.toml

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

{rxnn-0.1.60 → rxnn-0.1.61}/src/rxnn/transformers/models.py

@@ -72,11 +72,11 @@ class ReactiveTransformerDecoder(ReactiveTransformerBase):
         # Process shared layers
         if self.shared_layers is not None:
             for i in range(self.num_shared_layers):
-                layer_stm = self.stm(i)
+                layer_stm = self.stm(i).expand(x.size(0), -1, -1)
                 x = self.shared_layers[i](x, layer_stm, mask=mask)
         # Process own layers
         for i in range(self.num_own_layers):
-            layer_stm = self.stm(i)
+            layer_stm = self.stm(i).expand(x.size(0), -1, -1)
             x = self.layers[i](x, layer_stm, mask=mask)
         return self.head(x)
 
@@ -93,12 +93,12 @@ class ReactiveTransformerEncoder(ReactiveTransformerBase):
         # Process shared layers
         if self.shared_layers is not None:
             for i in range(self.num_shared_layers):
-                layer_stm = self.stm(i)
+                layer_stm = self.stm(i).expand(x.size(0), -1, -1)
                 x = self.shared_layers[i](x, layer_stm, mask=attention_mask)
                 hidden_states.append(x)
         # Process own layers
         for i in range(self.num_own_layers):
-            layer_stm = self.stm(i)
+            layer_stm = self.stm(i).expand(x.size(0), -1, -1)
             x = self.layers[i](x, layer_stm, mask=attention_mask)
             hidden_states.append(x)
         return x, torch.stack(hidden_states)