rxnn 0.1.63__tar.gz → 0.1.64__tar.gz

{rxnn-0.1.63 → rxnn-0.1.64}/PKG-INFO

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

{rxnn-0.1.63 → rxnn-0.1.64}/pyproject.toml

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

{rxnn-0.1.63 → rxnn-0.1.64}/src/rxnn/rxt/models.py

@@ -31,9 +31,11 @@ class RxTAlphaComponentConfig(TypedDict):
     moe_top_k: int
     self_att_type: str
     cross_att_type: str
-    att_num_experts: int
-    att_num_query_experts: int
-    att_num_query_groups: int
+    att_experts: int
+    att_query_experts: int
+    att_query_groups: int
+    cross_att_groups: int
+    cross_att_query_groups: int
 
 
 class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
@@ -61,9 +63,11 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
             moe_top_k: int = 1,
             self_att_type: str = 'gqa',
             cross_att_type: str = 'mqa',
-            att_num_experts: int = None,
-            att_num_query_experts: int = None,
-            att_num_query_groups: int = None,
+            att_experts: int = None,
+            att_query_experts: int = None,
+            att_query_groups: int = None,
+            cross_att_groups: int = None,
+            cross_att_query_groups: int = None,
             **kwargs
     ):
         super(RxTAlphaComponentBase, self).__init__(**kwargs)
@@ -86,20 +90,20 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
         else:
             att_init = lambda: init_experimental_attention(embed_dim, att_heads, self_att_type, att_groups, rope=rope,
                                                            use_flash_attention=use_flash_attention, dropout=att_dropout,
-                                                           max_seq_len=seq_len, is_causal=is_causal, num_experts=att_num_experts,
-                                                           num_query_experts=att_num_query_experts,
-                                                           num_query_groups=att_num_query_groups)
+                                                           max_seq_len=seq_len, is_causal=is_causal, num_experts=att_experts,
+                                                           num_query_experts=att_query_experts,
+                                                           num_query_groups=att_query_groups)
 
         if cross_att_type in ['mha', 'gqa', 'mqa']:
             cross_att_init = lambda: init_attention(embed_dim, att_heads, cross_att_type, att_groups, rope=rope,
                                               use_flash_attention=use_flash_attention, dropout=att_dropout,
                                               max_seq_len=seq_len, is_causal=is_causal)
         else:
-            cross_att_init = lambda: init_experimental_attention(embed_dim, att_heads, cross_att_type, att_groups, rope=rope,
+            cross_att_init = lambda: init_experimental_attention(embed_dim, att_heads, cross_att_type, cross_att_groups or att_groups, rope=rope,
                                                            use_flash_attention=use_flash_attention, dropout=att_dropout,
-                                                           max_seq_len=seq_len, is_causal=is_causal, num_experts=att_num_experts,
-                                                           num_query_experts=att_num_query_experts,
-                                                           num_query_groups=att_num_query_groups)
+                                                           max_seq_len=seq_len, is_causal=is_causal, num_experts=att_experts,
+                                                           num_query_experts=att_query_experts,
+                                                           num_query_groups=cross_att_query_groups or att_query_groups)
 
         layers = nn.ModuleList([
             ReactiveTransformerLayer(