rxnn 0.1.52__py3-none-any.whl → 0.1.53__py3-none-any.whl

rxnn/transformers/moe.py

@@ -23,20 +23,21 @@ class MoeRouter(nn.Module):
 
     def calculate_aux_loss(self, top_k_indices: torch.Tensor, probs: torch.Tensor) -> torch.Tensor:
         # Get shapes
-        B, S, K = top_k_indices.shape  # Batch, Sequence length, Top-K
+        T, K = top_k_indices.shape  # Batch, Sequence length, Top-K
 
         # 1. Compute expert selection mask (one-hot encoded)
         expert_mask = F.one_hot(top_k_indices, self.num_experts).float()  # (B, S, K, E)
 
         # 2. Total number of times each expert is selected
-        expert_usage = expert_mask.sum(dim=(0, 1, 2))  # (E,)
+        expert_usage = expert_mask.sum(dim=(0, 1))  # (E,)
 
         # 3. Fraction of tokens assigned to each expert
-        total_tokens = B * S * K
-        fraction_expert = expert_usage / total_tokens  # (E,)
+        total_selections = T * K
+        fraction_expert = expert_usage / (total_selections + 1e-6)  # (E,)
 
         # 4. Sum of probabilities for each expert's selected tokens
-        sum_probs = (probs.unsqueeze(-1) * expert_mask).sum(dim=(0, 1, 2))  # (E,)
+        probs_expanded = probs.unsqueeze(1).expand(-1, K, -1)  # (B_K, K, E)
+        sum_probs = (probs_expanded * expert_mask).sum(dim=(0, 1))
 
         # 5. Average probability per expert (avoid division by zero)
         avg_probs = sum_probs / expert_usage.clamp(min=1e-6)  # (E,)

{rxnn-0.1.52.dist-info → rxnn-0.1.53.dist-info}/METADATA

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

{rxnn-0.1.52.dist-info → rxnn-0.1.53.dist-info}/RECORD

@@ -21,11 +21,11 @@ rxnn/transformers/ff.py,sha256=jJnuBDsnnX5uYC_WZH8cXAYrMnz0P-iX7MwcPivjRtI,2533
 rxnn/transformers/layers.py,sha256=OX8CsFY9A7uqH1SLwyexR_5BNlwheYrJHCGXjF8Q7HU,7186
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
 rxnn/transformers/models.py,sha256=QFzBrOR7tDp9d_T0HoIukBMfEbLxsCictV5p3e2ilxg,7552
-rxnn/transformers/moe.py,sha256=88-w4cQhYNcebdq4zBsdkaoFa4VxJi1LFXDKAAkfVLk,5791
+rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 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.52.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.1.52.dist-info/METADATA,sha256=aae9Bt0SpsDgugeHY-7Bi6SN3wWhXneD3Kbz1NMtxJo,16627
-rxnn-0.1.52.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-rxnn-0.1.52.dist-info/RECORD,,
+rxnn-0.1.53.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.1.53.dist-info/METADATA,sha256=fwifQgrfowidD4AR5Pl952-P42HwtA3Tuq6d5RAM6QI,16627
+rxnn-0.1.53.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+rxnn-0.1.53.dist-info/RECORD,,