rxnn 0.1.21__tar.gz → 0.1.22__tar.gz

{rxnn-0.1.21 → rxnn-0.1.22}/PKG-INFO

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

{rxnn-0.1.21 → rxnn-0.1.22}/pyproject.toml

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

{rxnn-0.1.21 → rxnn-0.1.22}/src/rxnn/experimental/attention.py

@@ -229,11 +229,12 @@ class DeepMoeAttention(GroupedMoeAttention):
         weights_q = weights_q_flat.view(B, T, -1)
         indices_q = indices_q_flat.view(B, T, -1)
         q = self._process_grouped_experts(query, self.wq, self.bq, weights_q, indices_q)
+        print('processed q', q.size())
         q = q.permute(0, 2, 1, 3).reshape(B, self.num_query_groups, T, -1)
 
         # Expand query groups to match head count
         group_heads = self.num_heads // self.num_query_groups
-        q = q.unsqueeze(2).expand(-1, -1, group_heads, -1, -1).flatten(1, 2).transpose(1, 2)
+        q = q.unsqueeze(2).expand(-1, -1, group_heads, -1, -1).flatten(1, 2)
 
         # Key/Value processing
         return super()._forward_qkv(q, key, value, b, t, d, skip_query_processing=True)