rxnn 0.1.30__tar.gz → 0.1.31__tar.gz

{rxnn-0.1.30 → rxnn-0.1.31}/PKG-INFO

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

{rxnn-0.1.30 → rxnn-0.1.31}/pyproject.toml

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

{rxnn-0.1.30 → rxnn-0.1.31}/src/rxnn/experimental/attention.py

@@ -337,16 +337,26 @@ class GroupedMoeAttentionVectorized(GroupedQueryAttention):
         k_all = k_all.view(B, S, self.num_experts, -1)
         v_all = v_all.view(B, S, self.num_experts, -1)
 
+        print('k_all', k_all.size())
+        print('v_all', v_all.size())
+
+
         # Gather top-k experts and weights
         # Expand indices to [B, S, num_groups, head_dim]
         expanded_indices = indices.unsqueeze(-1).expand(-1, -1, -1, k_all.size(-1))
         selected_k = torch.gather(k_all, 2, expanded_indices)  # [B, S, num_groups, head_dim]
         selected_v = torch.gather(v_all, 2, expanded_indices)
 
+        print('selected_k', selected_k.size())
+        print('selected_v', selected_v.size())
+
         # Weighted sum
         weighted_k = (selected_k * weights).sum(dim=2)  # [B, S, head_dim]
         weighted_v = (selected_v * weights).sum(dim=2)
 
+        print('weighted_k', weighted_k.size())
+        print('weighted_v', weighted_v.size())
+
         # Reshape to GQA format
         k = weighted_k.view(B, S, self.num_groups, -1).permute(0, 2, 1, 3)  # [B, G, S, head_dim]
         v = weighted_v.view(B, S, self.num_groups, -1).permute(0, 2, 1, 3)