rxnn 0.1.20__tar.gz → 0.1.21__tar.gz

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

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

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

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

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

@@ -125,6 +125,9 @@ class GroupedMoeAttention(GroupedQueryAttention):
         k = self._process_grouped_experts(key, self.wk, self.bk, weights_k, indices_k)
         v = self._process_grouped_experts(value, self.wv, self.bv, weights_k, indices_k)
 
+        print('processed k', k.size())
+        print('processed v', v.size())
+
         # Expand to GQA format
         k = k.permute(0, 2, 1, 3).reshape(B, self.num_groups, S, -1)
         v = v.permute(0, 2, 1, 3).reshape(B, self.num_groups, S, -1)
@@ -138,6 +141,10 @@ class GroupedMoeAttention(GroupedQueryAttention):
             k = k.flatten(start_dim=1, end_dim=2)  # (B, H, S, head_dim)
             v = v.flatten(start_dim=1, end_dim=2)  # (B, H, S, head_dim)
 
+        print('q', q.size())
+        print('k', k.size())
+        print('v', v.size())
+
         return q, k, v