rxnn 0.1.27__py3-none-any.whl → 0.1.29__py3-none-any.whl

rxnn/experimental/attention.py

@@ -330,8 +330,8 @@ class GroupedMoeAttentionVectorized(GroupedQueryAttention):
 
         # Compute all experts' projections
         # Shape: (B*S, num_experts, head_dim)
-        k_all = torch.einsum('be,ehd->beh', key_flat, self.wk)  # [B*S, num_experts, head_dim]
-        v_all = torch.einsum('be,ehd->beh', value.view(B*S, D), self.wv)
+        k_all = torch.einsum('bd,edh->beh', key_flat, self.wk)  # [B*S, num_experts, head_dim]
+        v_all = torch.einsum('bd,edh->beh', value.view(B*S, D), self.wv)
 
         # Reshape to [B, S, num_experts, head_dim]
         k_all = k_all.view(B, S, self.num_experts, -1)
@@ -434,12 +434,13 @@ class DeepMoeAttentionVectorized(GroupedMoeAttentionVectorized):
 
     def _init_out(self, embed_dim: int):
         """Initialize output projection"""
-        self.out_proj = nn.Linear(embed_dim // (self.num_heads // self.num_groups), embed_dim)
+        out_hidden_dim = embed_dim // self.num_heads * self.num_query_groups
+        self.out_proj = nn.Linear(out_hidden_dim, embed_dim)
 
     def _transpose_output(self, attn_output: torch.Tensor, b: int, t: int, d: int):
         """Transpose attention output back to (B, T, D) shape"""
-        hidden_dim = d // self.num_heads * self.num_query_groups
-        return attn_output.transpose(1, 2).contiguous().view(b, t, hidden_dim)
+        out_hidden_dim = d // self.num_heads * self.num_query_groups
+        return attn_output.transpose(1, 2).contiguous().view(b, t, out_hidden_dim)
 
     def _forward_qkv(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, b: int, t: int, d: int):
         B, T, D = query.shape
@@ -449,7 +450,7 @@ class DeepMoeAttentionVectorized(GroupedMoeAttentionVectorized):
         indices_q = indices_q.view(B, T, self.num_query_groups)
 
         # Compute all query experts
-        q_all = torch.einsum('be,ehd->beh', query_flat, self.wq)  # [B*T, num_query_experts, head_dim]
+        q_all = torch.einsum('bd,edh->beh', query_flat, self.wq)  # [B*T, num_query_experts, head_dim]
         q_all = q_all.view(B, T, self.num_query_experts, -1)
 
         # Gather top-k experts

{rxnn-0.1.27.dist-info → rxnn-0.1.29.dist-info}/METADATA

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

{rxnn-0.1.27.dist-info → rxnn-0.1.29.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=csasMRxL4nq2dS7pc9WdS4bvCB70ZVgsR7LTHV2jEJ0,29388
+rxnn/experimental/attention.py,sha256=FLIyfRWZC1Ppfx9u8qUcNcfMSscuibhKEQ8zCJHXcfk,29439
 rxnn/experimental/models.py,sha256=QEuFBB9iEg5AbKQLwGJkAwPjMfaVeTqazhKDWPRkm7o,4598
 rxnn/experimental/moe.py,sha256=PhiaNr3FwR2Zv2a0tfj6sfZ4iyhLo3Jyp2DwXq19qZQ,7935
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -25,7 +25,7 @@ rxnn/transformers/moe.py,sha256=6Cffyo0QjmEWc4rK1ncOmLRCQbY0OpQJ4D7xH_4nTN4,4738
 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.27.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.27.dist-info/METADATA,sha256=XjcqSdhjTRsCvvP-o981Ihp4k5PFRCQUSLVsPZ_NVPw,16627
-rxnn-0.1.27.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.27.dist-info/RECORD,,
+rxnn-0.1.29.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.29.dist-info/METADATA,sha256=u3tlE_n8tZsqgW44vPssCmmJ03wHKM57tGS1xXI3HcQ,16627
+rxnn-0.1.29.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.29.dist-info/RECORD,,