rxnn 0.1.41__py3-none-any.whl → 0.1.42__py3-none-any.whl

rxnn/experimental/attention.py

@@ -128,8 +128,8 @@ class GroupedMoeAttention(GroupedQueryAttention):
         print(selected_v.size(), selected_v.dtype)
 
         # Weighted
-        weighted_k = selected_k * weights # [B, S, num_groups, head_dim]
-        weighted_v = selected_v * weights # [B, S, num_groups, head_dim]
+        weighted_k = (selected_k * weights).to(selected_k.device, dtype=selected_k.dtype) # [B, S, num_groups, head_dim]
+        weighted_v = (selected_v * weights).to(selected_k.device, dtype=selected_k.dtype)  # [B, S, num_groups, head_dim]
 
         print('weighted')
         print(weighted_k.size(), weighted_k.dtype)
@@ -253,7 +253,7 @@ class DeepMoeAttention(GroupedMoeAttention):
         selected_q = torch.gather(q_all, 2, expanded_indices)  # [B, T, num_query_groups, head_dim]
 
         # Weighted sum
-        q = selected_q * weights_q  # [B, T, num_query_groups, head_dim]
+        q = (selected_q * weights_q).to(selected_q.device, dtype=selected_q.dtype)  # [B, T, num_query_groups, head_dim]
         q = q.view(B, T, self.num_query_groups, -1).permute(0, 2, 1, 3)  # [B, num_query_groups, T, head_dim]
 
         return super()._forward_qkv(q, key, value, b, t, d, skip_query_processing=True)

{rxnn-0.1.41.dist-info → rxnn-0.1.42.dist-info}/METADATA

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

{rxnn-0.1.41.dist-info → rxnn-0.1.42.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=tvyjyBWLh9Bcu3MXlo3yx1hA48La8Zks4wWYtcnOYc8,30365
+rxnn/experimental/attention.py,sha256=mPDLI5lwujNTELdnVXDuIpagoQqHDP1GG6-ObCyM-Hw,30510
 rxnn/experimental/models.py,sha256=IzUVc5s-cA__8jsG2mVvzUDmzPRcfBcI5btaOjnPYhA,4598
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 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.41.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.41.dist-info/METADATA,sha256=B9AeelUbhWbH1uYB17WIEz4gjo_1c0V5jTM5S2QgulE,16627
-rxnn-0.1.41.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.41.dist-info/RECORD,,
+rxnn-0.1.42.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.42.dist-info/METADATA,sha256=38UxLA25RpEi1-how5kIxMtlyJSrqZWZNzw0sgtnoDs,16627
+rxnn-0.1.42.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.42.dist-info/RECORD,,