rxnn 0.1.59__tar.gz → 0.1.61__tar.gz

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

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

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

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

{rxnn-0.1.59 → rxnn-0.1.61}/src/rxnn/transformers/attention.py

@@ -137,6 +137,9 @@ class MultiHeadAttention(nn.Module):
         b, t, d = query.size()
         q, k, v = self._forward_qkv(query, key, value, b, t, d)
         if not self.rel_embed:
+            print('q', q.size())
+            print('k', k.size())
+            print('v', v.size())
             q, k = self._apply_rope(q, k)
             attn_output = self._calculate_attention(q, k, v, b, t, d, mask=mask)
         else:

{rxnn-0.1.59 → rxnn-0.1.61}/src/rxnn/transformers/layers.py

@@ -86,6 +86,7 @@ class ReactiveTransformerLayer(nn.Module):
         residual = x
         if not self.use_post_norm:
             x = self.norm1(x)
+        print('self x', x.size())
         x = self.attention(x, x, x, mask=mask)
         x = residual + x
         if self.use_post_norm:
@@ -94,6 +95,8 @@ class ReactiveTransformerLayer(nn.Module):
         residual = x
         if not self.use_post_norm:
             x = self.norm2(x)
+        print('x', x.size())
+        print('STM', stm.size())
         x = self.memory_cross_attention(x, stm, stm)
         x = residual + x
         if self.use_post_norm:

{rxnn-0.1.59 → 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)