rxnn 0.2.58__py3-none-any.whl → 0.2.59__py3-none-any.whl

rxnn/memory/attention.py

@@ -48,7 +48,9 @@ class StmMemoryAttention(nn.Module):
             return layer_gate * new_layer_stm + (1 - layer_gate) * layer_stm
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
-        mem_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool() if attention_mask else None
+        if attention_mask is not None:
+            print(attention_mask.size())
+            attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
         new_stm = torch.zeros_like(self.stm.memory)
         for i in range(self.num_layers):
             layer_stm = self.stm(i)
@@ -57,7 +59,7 @@ class StmMemoryAttention(nn.Module):
                 layer_stm = layer_stm.expand(x.size(0), -1, -1)
             encoded_layer_data = x[i]
             normalized_layer_stm = self.memory_norm_layers[i](layer_stm)
-            new_layer_stm = self.attention_layers[i](normalized_layer_stm, encoded_layer_data, encoded_layer_data, mask=mem_mask)
+            new_layer_stm = self.attention_layers[i](normalized_layer_stm, encoded_layer_data, encoded_layer_data, mask=attention_mask)
             if self.use_gated_residual:
                 new_stm[i] = self._residual_gate(self.gate[i], layer_stm, new_layer_stm) # gated residual
             else:

rxnn/rxt/models.py

@@ -103,13 +103,13 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
         if cross_att_type in ['mha', 'gqa', 'mqa']:
             cross_att_init = lambda: init_attention(embed_dim, att_heads, cross_att_type, att_groups, rope=rope,
                                                     use_flash_attention=use_flash_attention, dropout=att_dropout,
-                                                    max_seq_len=seq_len, is_causal=is_causal, rope_only_for_query=True)
+                                                    max_seq_len=seq_len, is_causal=False, rope_only_for_query=True)
         else:
             cross_att_init = lambda: init_experimental_attention(embed_dim, att_heads, cross_att_type,
                                                                  cross_att_groups or att_groups, rope=rope,
                                                                  use_flash_attention=use_flash_attention,
                                                                  dropout=att_dropout,
-                                                                 max_seq_len=seq_len, is_causal=is_causal,
+                                                                 max_seq_len=seq_len, is_causal=False,
                                                                  num_experts=att_experts,
                                                                  num_query_experts=att_query_experts,
                                                                  num_query_groups=cross_att_query_groups or att_query_groups,

rxnn/transformers/models.py

@@ -108,6 +108,7 @@ class ReactiveTransformerEncoder(ReactiveTransformerBase):
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> tuple[torch.Tensor, torch.Tensor]:
         x = super().forward(x)  # apply embeddings
         if attention_mask is not None:
+            print(attention_mask.size())
             attention_mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
 
         hidden_states = []
@@ -213,6 +214,7 @@ class ClassicTransformerDecoder(ClassicTransformerBase):
             if attention_mask is not None:
                 mask &= attention_mask.unsqueeze(1).unsqueeze(1).bool()
         elif attention_mask is not None:
+            print(attention_mask.size())
             mask = attention_mask.unsqueeze(1).unsqueeze(1).bool()
         else:
             mask = None

{rxnn-0.2.58.dist-info → rxnn-0.2.59.dist-info}/METADATA

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

{rxnn-0.2.58.dist-info → rxnn-0.2.59.dist-info}/RECORD

@@ -5,11 +5,11 @@ rxnn/experimental/attention.py,sha256=jlNS82INjycNEfmk3HtkIacUvT_ELhaCO2g-kZTvhX
 rxnn/experimental/models.py,sha256=oJWd56LUsLc9S8eCZw-ShvuWjoQxj4C9GitbohlQ0ok,5139
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/memory/attention.py,sha256=N78kzcqXXfcq5v43LUsLheVUs4JrKPeQissl-_XKXdk,3241
+rxnn/memory/attention.py,sha256=Mt9LC0DK5asjN4pTqWOLDAoGzAgg3FhQoUw6baDJ3NI,3309
 rxnn/memory/norm.py,sha256=cVjjhCLqR5K6-321SP_ObG17y-ddlcTJeCTXvW4vpk0,6675
 rxnn/memory/stm.py,sha256=jv57gsH9XW19sLbxpRDqsp1yfsii_4Ef4Ncr_ztk-i4,3937
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/rxt/models.py,sha256=4MbCL4xGY3ceewZQmopjmwAyLQS92L6KLOPqaW7-Fho,14673
+rxnn/rxt/models.py,sha256=jd1UVBUWJzWdw7Rjcvo9k5BXCJriQ0khuVszqEyfD7M,14665
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=CqaArEZYOdH64nmKfx28U3GI46TzO4oNkjf_hrF23Cw,11835
 rxnn/training/bml.py,sha256=hw6gLpLkGvqLzxIvBg4MvCc5r8cHpEm2RDyh7nH6CtE,16914
@@ -28,12 +28,12 @@ rxnn/transformers/attention.py,sha256=KRnKT6XUqAXElxV9y72mSpdTeiMgCKCCLqqxCFNTHm
 rxnn/transformers/ff.py,sha256=WDjO-H9XWInoWnUnxiseIH6Kx5GlHP0zGJygwhcb1gc,2589
 rxnn/transformers/layers.py,sha256=OlbqD5kKygn5WZziLbU3jZjhr8hBrxLpqlCjJ_BNCW0,8119
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
-rxnn/transformers/models.py,sha256=7ypPNFFnacdZjvaLVue1KR2PmMSdVYsbCMQSunXDL70,10720
+rxnn/transformers/models.py,sha256=bHR5gy74aV20Sl0371vVCvb3Z2pcqYTdOSnrLSkAIiI,10802
 rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=t6iiQTdLQ0TakUWnnhKkb5DKF2F_9-thXHBydDF3fxg,17389
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.58.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.58.dist-info/METADATA,sha256=_NKVaBMYEbJadkBcUvDX1UprMlUu92JqgnBYx7R1J1c,25997
-rxnn-0.2.58.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.58.dist-info/RECORD,,
+rxnn-0.2.59.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.59.dist-info/METADATA,sha256=yj3LmcYzHTe7O5Y3LuDkCbjoTcY01WRQxaPQ9G4a8GA,25997
+rxnn-0.2.59.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.59.dist-info/RECORD,,