rxnn 0.1.40__tar.gz → 0.1.41__tar.gz

{rxnn-0.1.40 → rxnn-0.1.41}/PKG-INFO

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

{rxnn-0.1.40 → rxnn-0.1.41}/pyproject.toml

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

{rxnn-0.1.40 → rxnn-0.1.41}/src/rxnn/experimental/attention.py

@@ -90,6 +90,7 @@ class GroupedMoeAttention(GroupedQueryAttention):
 
         # Key/Value MoE routing
         B, S, D = key.shape
+        print('key/value type', key.dtype, value.dtype)
         key_flat = key.reshape(-1, D)
         print('key flat type', key_flat.dtype)
         weights, indices = self.router(key_flat)  # (B*S, num_groups), (B*S, num_groups)

{rxnn-0.1.40 → rxnn-0.1.41}/src/rxnn/transformers/attention.py

@@ -137,6 +137,10 @@ class MultiHeadAttention(nn.Module):
         return self._calculate_output(attn_weights, v, b, t, d)
 
     def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: torch.Tensor = None):
+        print('MHA forward')
+        print(query.size(), query.dtype)
+        print(key.size(), key.dtype)
+        print(value.size(), value.dtype)
         b, t, d = query.size()
         q, k, v = self._forward_qkv(query, key, value, b, t, d)
         if not self.rel_embed: