rxnn 0.1.17__py3-none-any.whl → 0.1.19__py3-none-any.whl

rxnn/experimental/attention.py

@@ -100,7 +100,7 @@ class GroupedMoeAttention(GroupedQueryAttention):
 
             # Get positions where expert e is used in any group
             x_slice = x_flat[token_mask]
-            proj = F.linear(x_slice, w[e], b[e] if b is not None else None)
+            proj = F.linear(x_slice, w[e].t(), b[e] if b is not None else None)
 
             # Find which groups use this expert for selected tokens
             group_mask = (indices_flat[token_mask] == e)
@@ -118,6 +118,7 @@ class GroupedMoeAttention(GroupedQueryAttention):
         # Key/Value processing
         B, S, D = key.shape
         key_flat = key.view(-1, D)
+        print('key_flat: ', key_flat.shape)
         weights_k_flat, indices_k_flat = self.router(key_flat)
         # Reshape back to original dimensions
         weights_k = weights_k_flat.view(B, S, -1)

rxnn/transformers/moe.py

@@ -16,22 +16,26 @@ class MoeRouter(nn.Module):
 
     def calculate_aux_loss(self, top_k_indices: torch.Tensor, probs: torch.Tensor) -> torch.Tensor:
         expert_mask = F.one_hot(top_k_indices, self.num_experts).float()
+        print('expert mask: ', expert_mask.shape)
         expert_usage = expert_mask.sum(dim=0).mean(dim=0)
+        print('expert usage: ', expert_usage.shape)
         mean_probs = probs.mean(dim=0)
+        print('mean probs: ', mean_probs.shape)
         return (expert_usage * mean_probs).sum() * self.num_experts
 
 
     def forward(self, x: torch.Tensor):
         # Input shape: [batch*seq_len, embed_dim]
         logits = self.gate(x)
+        print('router logits: ', logits.shape)
         probs = F.softmax(logits, dim=-1)
-
+        print('router probs: ', probs.shape)
         # Get top-k experts for each token
         top_k_weights, top_k_indices = probs.topk(self.top_k, dim=-1)
 
         # Normalize weights (sum to 1 for each token)
         top_k_weights = top_k_weights / (top_k_weights.sum(dim=-1, keepdim=True) + 1e-9)
-
+        print('top k: ', top_k_weights.shape, top_k_indices.shape)
         # Load Balance Loss
         self.aux_loss = self.calculate_aux_loss(top_k_indices, probs)
 

{rxnn-0.1.17.dist-info → rxnn-0.1.19.dist-info}/METADATA

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

{rxnn-0.1.17.dist-info → rxnn-0.1.19.dist-info}/RECORD

@@ -1,6 +1,6 @@
 rxnn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/experimental/attention.py,sha256=wjHrxfov3Ybg3iou8FlQtFvxNuHdcs_A7a6FTloosgA,32056
+rxnn/experimental/attention.py,sha256=nvYtC6BdJQ8VUcNc_co2Fe2at7TBzA4OOIfG2tWWVCk,32104
 rxnn/experimental/models.py,sha256=-XkEHsyT8iNAjhZbgC7N_5nzP4ENVJLwxSoLHgMfA0I,4668
 rxnn/experimental/moe.py,sha256=PhiaNr3FwR2Zv2a0tfj6sfZ4iyhLo3Jyp2DwXq19qZQ,7935
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -21,11 +21,11 @@ rxnn/transformers/ff.py,sha256=jJnuBDsnnX5uYC_WZH8cXAYrMnz0P-iX7MwcPivjRtI,2533
 rxnn/transformers/layers.py,sha256=HhIiykmrBgdsV4AbMQXr9t0cSo4gSIeN0dPtc8mDyOo,5629
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
 rxnn/transformers/models.py,sha256=w-zB_8QB9-Fae-GkGgmVDNY-Ts_0gBeWcevpl9qzZVM,7169
-rxnn/transformers/moe.py,sha256=FeaQR7hTX1dE74YdMOcuyZHSkGiV_0JwF8fw-GnfNOQ,4741
+rxnn/transformers/moe.py,sha256=gJ-jXKtc01xcBayaYchRZy7imFGnvwVfUflXvFiKjKU,5048
 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.17.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.17.dist-info/METADATA,sha256=wId6o7JCcBjRD1plWzgJRmFAY5VlHN7-FIVySeVDqx8,16627
-rxnn-0.1.17.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.17.dist-info/RECORD,,
+rxnn-0.1.19.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.19.dist-info/METADATA,sha256=4ul6X1SOT2bzHCxK88SjcYc0-1zy8YAKPCoMtZ2dKrY,16627
+rxnn-0.1.19.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.19.dist-info/RECORD,,