rxnn 0.1.21__py3-none-any.whl → 0.1.23__py3-none-any.whl

rxnn/experimental/attention.py

@@ -125,9 +125,6 @@ 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)
@@ -141,10 +138,6 @@ 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
 
 
@@ -229,12 +222,8 @@ class DeepMoeAttention(GroupedMoeAttention):
         weights_q = weights_q_flat.view(B, T, -1)
         indices_q = indices_q_flat.view(B, T, -1)
         q = self._process_grouped_experts(query, self.wq, self.bq, weights_q, indices_q)
-        q = q.permute(0, 2, 1, 3).reshape(B, self.num_query_groups, T, -1)
-
-        # Expand query groups to match head count
-        group_heads = self.num_heads // self.num_query_groups
-        q = q.unsqueeze(2).expand(-1, -1, group_heads, -1, -1).flatten(1, 2).transpose(1, 2)
 
+        q = q.permute(0, 2, 1, 3).reshape(B, self.num_query_groups, T, -1)
         # Key/Value processing
         return super()._forward_qkv(q, key, value, b, t, d, skip_query_processing=True)
 

{rxnn-0.1.21.dist-info → rxnn-0.1.23.dist-info}/METADATA

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

{rxnn-0.1.21.dist-info → rxnn-0.1.23.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=6qD3QCpkQHsIaNktjcQrRitQgQ-WkRUVtSFgEDfYGbA,32340
+rxnn/experimental/attention.py,sha256=gMEcFJHGOkz8R_s4dGEJB5cb2K3pbXZi4XBwyhEdB4s,31967
 rxnn/experimental/models.py,sha256=-XkEHsyT8iNAjhZbgC7N_5nzP4ENVJLwxSoLHgMfA0I,4668
 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=msspVdefdt2ekIN8aT-V8DolK4taESQL_NVsSGOepIs,4739
 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.21.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.21.dist-info/METADATA,sha256=gtCrs3sVTMXB9UNS1-qcJNIPzHNO8d7UaJlfviJNFEI,16627
-rxnn-0.1.21.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.21.dist-info/RECORD,,
+rxnn-0.1.23.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.23.dist-info/METADATA,sha256=rZSBuoIgf8jKB11LKgMg7U42Wx7VNT_4EU3FVyED2YQ,16627
+rxnn-0.1.23.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.23.dist-info/RECORD,,