rxnn 0.1.60__py3-none-any.whl → 0.1.62__py3-none-any.whl

rxnn/transformers/attention.py

@@ -137,9 +137,6 @@ 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/transformers/layers.py

@@ -86,7 +86,6 @@ 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:
@@ -95,8 +94,6 @@ 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/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)

{rxnn-0.1.60.dist-info → rxnn-0.1.62.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.1.60
+Version: 0.1.62
 Summary: RxNN: Reactive Neural Networks Platform
 License: Apache-2.0
 Keywords: deep-learning,ai,machine-learning
@@ -14,7 +14,6 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: datasets (>=3.5.0,<4.0.0)
-Requires-Dist: flash-attention (>=1.0.0,<2.0.0)
 Requires-Dist: huggingface-hub (>=0.30.0,<0.31.0)
 Requires-Dist: tensorboard (>=2.19.0,<3.0.0)
 Requires-Dist: tokenizers (>=0.21.0,<0.22.0)

{rxnn-0.1.60.dist-info → rxnn-0.1.62.dist-info}/RECORD

@@ -16,16 +16,16 @@ rxnn/training/dataset.py,sha256=JQuWSUdT5AnsrG6M_EsewoU6uroVHhg4K715nbtDx8A,9643
 rxnn/training/scheduler.py,sha256=ow6oALzWjWQmHSpcJEjv6tg4g4CDMvr73TypxfcefMc,712
 rxnn/training/tokenizer.py,sha256=umaLByMBx_NMrQElA45HLm9gkuzyKWDTFaKVd-CjXl0,8344
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/transformers/attention.py,sha256=2pJoAqRB6AKHtA2lj158NzbHzXr0dCaKj1RKjMfHsOI,15794
+rxnn/transformers/attention.py,sha256=dC0UmC-_kjX8US6Sf0Fi5zw5kJ-P6orH3JDHeBB5gI8,15695
 rxnn/transformers/ff.py,sha256=jJnuBDsnnX5uYC_WZH8cXAYrMnz0P-iX7MwcPivjRtI,2533
-rxnn/transformers/layers.py,sha256=tSOwr_IwJcAjPLmyRJI-93wRjgHH6OPJj4QrdkEFYdc,7282
+rxnn/transformers/layers.py,sha256=OX8CsFY9A7uqH1SLwyexR_5BNlwheYrJHCGXjF8Q7HU,7186
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
-rxnn/transformers/models.py,sha256=QFzBrOR7tDp9d_T0HoIukBMfEbLxsCictV5p3e2ilxg,7552
+rxnn/transformers/models.py,sha256=xbnn3FTNZFhaqq9A0XEM12ie_WL_58pPeq0qFXIgve0,7656
 rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=2l38RS0Dini3f6Z3LUHr3XwWzg1UK7fO2C6wazWDAYU,4292
 rxnn/transformers/sampler.py,sha256=poWBpxg1iuK5gEJtxHkk5VVfS9V48hs2Olqdhy_Gw8c,6548
 rxnn/utils.py,sha256=d5U8i5ukovgDyqiycc2AoxObTz_eF_bgo2MKvdtJ98s,467
-rxnn-0.1.60.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.60.dist-info/METADATA,sha256=eTT7LMBzM_AiWqreCwois6pTV50-zb8tyMlgrz4Mq-g,16627
-rxnn-0.1.60.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.60.dist-info/RECORD,,
+rxnn-0.1.62.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.62.dist-info/METADATA,sha256=ku_ZxKtWSEaqJ5qeRmlAeKSk69upNr_CWRgoZel3yOo,16579
+rxnn-0.1.62.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.62.dist-info/RECORD,,