rxnn 0.1.52__tar.gz → 0.1.53__tar.gz

{rxnn-0.1.52 → rxnn-0.1.53}/PKG-INFO

@@ -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 → rxnn-0.1.53}/pyproject.toml

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

{rxnn-0.1.52 → rxnn-0.1.53}/src/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,)