rxnn 0.1.27__tar.gz → 0.1.28__tar.gz

{rxnn-0.1.27 → rxnn-0.1.28}/PKG-INFO

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

{rxnn-0.1.27 → rxnn-0.1.28}/pyproject.toml

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

{rxnn-0.1.27 → rxnn-0.1.28}/src/rxnn/experimental/attention.py

@@ -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