rxnn 0.2.24__py3-none-any.whl → 0.2.26__py3-none-any.whl

rxnn/memory/attention.py

@@ -35,7 +35,6 @@ class StmMemoryAttention(nn.Module):
             encoded_layer_data = x[i]
             normalized_layer_stm = self.memory_norm_layers[i](layer_stm)
             new_layer_stm = self.attention_layers[i](normalized_layer_stm, encoded_layer_data, encoded_layer_data, mask=mask)
-            # self.stm.update_layer(i, new_layer_stm + layer_stm)
             new_stm[i] = new_layer_stm + layer_stm # residual
         self.stm.update_all(new_stm)
         return self.stm.memory

rxnn/memory/norm.py

@@ -7,10 +7,11 @@ class AdaptivePositionalMemoryNorm(nn.Module):
         self,
         num_slots: int,
         dim: int,
-        decay: float = 0.99,
+        decay: float = 0.9,
         use_scale: bool = True,
         use_gate: bool = True,
-        init_gate: float = -4.0
+        init_gate: float = -2.0,
+        per_dim_scale: bool = False,
     ):
         super(AdaptivePositionalMemoryNorm, self).__init__()
         self.use_gate = use_gate
@@ -20,39 +21,38 @@ class AdaptivePositionalMemoryNorm(nn.Module):
         self.eps = 1e-6
 
         # Learnable parameters
-        self.scale = nn.Parameter(torch.ones(num_slots, 1, dim)) if use_scale else None
-        self.gate = nn.Parameter(torch.full((num_slots, 1, 1), init_gate)) if use_gate else None
+        scale_shape = (num_slots, 1) if not per_dim_scale else (dim,)
+        self.scale = nn.Parameter(torch.ones(*scale_shape)) if use_scale else None
+        self.gate = nn.Parameter(torch.full((num_slots, 1), init_gate)) if use_gate else None
 
         # EMA buffers
         self.register_buffer("ema_rms", torch.ones(num_slots, 1))
 
         # Initialize parameters
         if self.scale is not None:
-            nn.init.normal_(self.scale, mean=1.0, std=0.01)
+            nn.init.normal_(self.scale, mean=1.0, std=0.1)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # x shape: [batch_size, num_slots, dim]
-        batch_size = x.size(0)
-
         # Calculate current RMS per slot
-        current_rms = x.pow(2).mean(dim=-1, keepdim=True).sqrt()  # [batch, slots, 1]
-        slot_rms = current_rms.mean(dim=0)  # [slots, 1] (average over batch)
+        # x: [batch_size, num_slots, dim]
+        current_rms = x.pow(2).mean(dim=-1, keepdim=True).sqrt()  # [batch, num_slots, 1]
+        slot_rms = current_rms.mean(dim=0)  # [num_slots, 1] (average over batch)
 
         # Update EMA during training
         if self.training:
-            self.ema_rms = self.decay * self.ema_rms + (1 - self.decay) * slot_rms.detach()
+            self.ema_rms = self.decay * self.ema_rms + (1 - self.decay) * slot_rms.detach() # [num_slots, 1]
 
         # Normalize using EMA statistics
-        x_norm = x * torch.rsqrt(self.ema_rms + self.eps)
+        x_norm = x * torch.rsqrt(self.ema_rms + self.eps) # [batch_size, num_slots, dim] * [num_slots, 1]
 
         # Apply learned scale per slot
         if self.scale is not None:
-            x_norm = x_norm * self.scale
+            x_norm = x_norm * self.scale # [batch_size, num_slots, dim] * [num_slots, 1] or [dim]
 
         # Apply gating mechanism
         if self.use_gate:
-            gate = torch.sigmoid(self.gate)  # [slots, 1, 1]
-            return gate * x_norm + (1 - gate) * x
+            gate = torch.sigmoid(self.gate)  # [num_slots, 1]
+            return gate * x_norm + (1 - gate) * x # [batch_size, num_slots, dim] * [num_slots, 1]
 
         return x_norm
 
@@ -77,7 +77,7 @@ class AdaptiveRMSMemoryNorm(nn.Module):
         # x shape: [batch_size, num_slots, dim]
         if self.training and hasattr(self, 'ema_rms'):
             # Compute current RMS across all slots and batch (scalar)
-            current_rms = x.pow(2).mean(-1).mean().sqrt()
+            current_rms = x.pow(2).mean(dim=-1).mean().sqrt()
             self.ema_rms = self.ema_rms * self.decay + current_rms * (1 - self.decay)
             rms = self.ema_rms
         else:
@@ -150,24 +150,26 @@ class MemoryNormConfig(TypedDict):
     use_gate: bool
     init_gate: float
     init_scale: float
+    per_dim_scale: bool
 
 def init_memory_norm(
     norm_type: str,
     dim: int,
     num_slots: int = None,
-    decay: float = 0.99,
+    decay: float = 0.9,
     use_scale: bool = True,
     use_gate: bool = True,
-    init_gate: float = -4.0,
+    init_gate: float = -2.0,
     init_scale: float = 1.0,
+    per_dim_scale: bool = False,
 ) -> nn.Module:
-    assert norm_type in ["layer", "rms", "adaptive", "positional"]
-    if norm_type == "layer":
+    assert norm_type in ['layer', 'rms', 'adaptive', 'positional']
+    if norm_type == 'layer':
         return MemoryLayerNorm(dim, use_gate, init_scale, init_gate)
-    elif norm_type == "rms":
+    elif norm_type == 'rms':
         return SimpleRMSMemoryNorm(dim, use_gate, init_scale, init_gate)
-    elif norm_type == "adaptive":
+    elif norm_type == 'adaptive':
         return AdaptiveRMSMemoryNorm(dim, use_gate, decay, init_scale, init_gate)
-    elif norm_type == "positional":
-        return AdaptivePositionalMemoryNorm(num_slots, dim, decay, use_scale, use_gate, init_gate)
+    elif norm_type == 'positional':
+        return AdaptivePositionalMemoryNorm(num_slots, dim, decay, use_scale, use_gate, init_gate, per_dim_scale)
     return MemoryLayerNorm(dim, use_gate, init_scale, init_gate)

rxnn/rxt/models.py

@@ -13,6 +13,7 @@ from ..memory.attention import StmMemoryAttention
 from ..utils import get_model_size
 from ..experimental.attention import init_experimental_attention
 
+
 class RxTAlphaComponentConfig(TypedDict):
     num_layers: int
     vocab_size: int
@@ -76,8 +77,10 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
         assert ff_activation in ['relu', 'gelu',
                                  'swish', 'silu', 'linear',
                                  'sigmoid'], 'Feed-forward activation could be "relu", "gelu", "swish", "silu", "linear", "sigmoid".'
-        assert self_att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma', 'sqa'], 'Self-attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
-        assert cross_att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma', 'sqa'], 'Memory cross-attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
+        assert self_att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma',
+                                 'sqa'], 'Self-attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
+        assert cross_att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma',
+                                  'sqa'], 'Memory cross-attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
 
         embedding = nn.Embedding(vocab_size, embed_dim)
         rope = RotaryPositionalEmbedding(embed_dim // att_heads, seq_len)
@@ -92,20 +95,25 @@ 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_experts,
+                                                           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, rope_only_for_query=True)
+                                                    use_flash_attention=use_flash_attention, dropout=att_dropout,
+                                                    max_seq_len=seq_len, is_causal=is_causal, rope_only_for_query=True)
         else:
-            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_experts,
-                                                           num_query_experts=att_query_experts,
-                                                           num_query_groups=cross_att_query_groups or att_query_groups, rope_only_for_query=True)
+            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_experts,
+                                                                 num_query_experts=att_query_experts,
+                                                                 num_query_groups=cross_att_query_groups or att_query_groups,
+                                                                 rope_only_for_query=True)
 
         layers = nn.ModuleList([
             ReactiveTransformerLayer(
@@ -137,6 +145,12 @@ class RxTAlphaComponentBase(nn.Module, PyTorchModelHubMixin):
     def load_shared_memory(self, stm: ShortTermMemory):
         self.model.stm = stm
 
+    def memory_parameters(self) -> list[nn.Parameter]:
+        return self.model.memory_parameters()
+
+    def not_memory_parameters(self) -> list[nn.Parameter]:
+        return self.model.not_memory_parameters()
+
     def freeze_without_memory(self, unfreeze_norms: bool = True):
         for param in self.model.parameters():
             param.requires_grad_(False)
@@ -211,20 +225,9 @@ class RxTAlphaDecoder(RxTAlphaComponentBase, pipeline_tag="text-generation", lic
         return self.model(x, attention_mask=attention_mask)
 
 
-def build_rxt_alpha_for_pretraining(
-        encoder_config: RxTAlphaComponentConfig,
-        decoder_config: RxTAlphaComponentConfig,
-) -> tuple[RxTAlphaEncoder, RxTAlphaDecoder]:
-    encoder = RxTAlphaEncoder(**encoder_config)
-    decoder = RxTAlphaDecoder(**decoder_config)
-
-    encoder.load_shared_memory(decoder.model.stm)
-    encoder.load_shared_embedding(decoder.model.embedding)
-
-    return encoder, decoder
-
 class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2.0"):
     """RxT-Alpha (Reactive Transformer) memory attention model"""
+
     def __init__(
             self,
             num_layers: int = 12,
@@ -234,17 +237,21 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
             stm_size: int = 1024,
             use_flash_attention: bool = False,
             att_dropout: float = 0.0,
-            norm_type: str = 'rms',
             att_groups: int = 1,
             att_type: str = 'sqa',
             att_experts: int = None,
             att_query_experts: int = None,
             att_query_groups: int = None,
+            norm_type: str = 'rms',
+            norm_init_gate: float = -2.0,
+            norm_per_dim_scale: bool = False,
+            norm_decay: float = 0.9,
             **kwargs,
     ):
         super(RxTAlphaMemoryAttention, self).__init__(**kwargs)
 
-        assert att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma', 'sqa'], 'Memory attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
+        assert att_type in ['mha', 'gqa', 'mqa', 'gma', 'dma',
+                            'sqa'], 'Memory attention type could be "mha", "gqa", "mqa", "gma", "dma", "sqa".'
 
         rope = RotaryPositionalEmbedding(embed_dim // att_heads, seq_len)
         stm = ShortTermMemory(num_layers, embed_dim, stm_size)
@@ -256,11 +263,14 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
         else:
             att_init = lambda: init_experimental_attention(embed_dim, att_heads, att_type, att_groups, rope=rope,
                                                            use_flash_attention=use_flash_attention, dropout=att_dropout,
-                                                           max_seq_len=seq_len, is_causal=False, num_experts=att_experts,
+                                                           max_seq_len=seq_len, is_causal=False,
+                                                           num_experts=att_experts,
                                                            num_query_experts=att_query_experts,
                                                            num_query_groups=att_query_groups, rope_only_for_keys=True)
 
-        memory_norm_layers = nn.ModuleList([init_memory_norm(norm_type, embed_dim, stm_size) for _ in range(num_layers)])
+        memory_norm_layers = nn.ModuleList([init_memory_norm(norm_type, embed_dim, stm_size, decay=norm_decay,
+                                                             init_gate=norm_init_gate, per_dim_scale=norm_per_dim_scale)
+                                            for _ in range(num_layers)])
         attention_layers = nn.ModuleList([att_init() for _ in range(num_layers)])
         self.model = StmMemoryAttention(stm, attention_layers, memory_norm_layers)
 
@@ -283,4 +293,3 @@ class RxTAlphaMemoryAttention(nn.Module, PyTorchModelHubMixin, license="apache-2
 
     def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
         return self.model(x, attention_mask=attention_mask)
-

rxnn/training/models.py

@@ -124,6 +124,19 @@ class MrlActorModel(nn.Module):
     def reset_memory(self):
         self.memory_attention.reset_memory()
 
+    def memory_parameters(self) -> list[nn.Parameter]:
+        return list(set(
+            self.encoder.memory_parameters() +
+            self.decoder.memory_parameters() +
+            self.memory_attention.parameters()
+        ))
+
+    def not_memory_parameters(self) -> list[nn.Parameter]:
+        return list(set(
+            self.encoder.not_memory_parameters() +
+            self.decoder.not_memory_parameters()
+        ))
+
     def unique_parameters(self):
         return list(set(
             list(self.encoder.parameters()) +

rxnn/training/mrl.py

@@ -17,6 +17,8 @@ from .models import MrlActorAction, MrlActorModel, MrlCriticModel
 
 class MrlConfig(TypedDict):
     lr: float
+    separate_memory_lr: Optional[bool]
+    memory_lr: Optional[float]
     critic_lr: float
     max_seq_len: int
     critic_max_len: int
@@ -42,7 +44,9 @@ class CurriculumConfig(TypedDict):
     random_resets_from: Optional[int]
     random_resets_ratio: Optional[float]
     reward_model: Optional[MrlRewardModel]
+    separate_memory_lr: Optional[bool]
     lr: Optional[float]
+    memory_lr: Optional[float]
     critic_lr: Optional[float]
     weight_decay: Optional[float]
     critic_weight_decay: Optional[float]
@@ -84,6 +88,7 @@ class MRLTrainer:
             use_amp: bool = False,
             dtype: torch.dtype = torch.float32,
             callbacks: list[MrlTrainerCallback] = None,
+
     ):
         """
         Trainer for Memory Reinforcement Learning (MRL) in Reactive Transformer.
@@ -123,15 +128,25 @@ class MRLTrainer:
         self.use_amp = use_amp
         self.dtype = dtype
 
-        self.base_optim_config = {
-            'lr': config.get('lr', 3e-4),
-            'critic_lr': config.get('critic_lr', 1e-4),
-            'weight_decay': config.get('weight_decay', 0.01),
-            'critic_weight_decay': config.get('critic_weight_decay', 0.01),
-        }
+        self.separate_memory_lr = config.get('separate_memory_lr', False)
+
+        if self.separate_memory_lr:
+            self.base_optim_config = {
+                'lr': (config.get('lr', 3e-4), config.get('memory_lr', 5e-4)),
+                'critic_lr': config.get('critic_lr', 1e-4),
+                'weight_decay': config.get('weight_decay', 0.01),
+                'critic_weight_decay': config.get('critic_weight_decay', 0.01),
+            }
+        else:
+            self.base_optim_config = {
+                'lr': config.get('lr', 3e-4),
+                'critic_lr': config.get('critic_lr', 1e-4),
+                'weight_decay': config.get('weight_decay', 0.01),
+                'critic_weight_decay': config.get('critic_weight_decay', 0.01),
+            }
 
         # Optimizers
-        self.optimizer, self.critic_optimizer = self._init_optimizers(**self.base_optim_config)
+        self.optimizer, self.critic_optimizer = self._init_optimizers(**self.base_optim_config, separate_memory_lr=self.separate_memory_lr)
 
         self.scaler = torch.amp.GradScaler() if self.use_amp else None
         self.critic_scaler = torch.amp.GradScaler() if self.use_amp else None
@@ -158,18 +173,28 @@ class MRLTrainer:
         self.global_epoch = 0
         self.global_epochs_count = 0
 
-    def _init_optimizers(self, lr: float, critic_lr: float, weight_decay: float, critic_weight_decay: float):
-        optimizer = torch.optim.AdamW(
-            self.actor.unique_parameters(),
-            lr=lr,
-            weight_decay=weight_decay,
-        )
+    def _init_optimizers(self, lr: Union[float, tuple[float, float]], critic_lr: float, weight_decay: float, critic_weight_decay: float, separate_memory_lr: bool = False) -> tuple[torch.optim.Optimizer, torch.optim.Optimizer]:
+        if separate_memory_lr:
+            rest_lr, memory_lr = lr
+            optimizer = torch.optim.AdamW([
+                { 'params': self.actor.not_memory_parameters(), 'lr': rest_lr },
+                { 'params': self.actor.memory_parameters(), 'lr': memory_lr },
+            ],
+                weight_decay=weight_decay,
+            )
+        else:
+            optimizer = torch.optim.AdamW(
+                self.actor.unique_parameters(),
+                lr=lr,
+                weight_decay=weight_decay,
+            )
 
         critic_optimizer = torch.optim.AdamW(
             self.critic.parameters(),
             lr=critic_lr,
             weight_decay=critic_weight_decay,
         )
+
         return optimizer, critic_optimizer
 
 
@@ -722,12 +747,13 @@ class MRLTrainer:
         self.strategy = config.get('strategy',
                                    MrlStrategy.MULTI_STEP_STRATEGY)  # MRL strategy for given curriculum stage
         self.reward = config.get('reward_model', self.shared_reward_model)  # MRL Reward Model for curriculum stage
-        if config['lr'] is not None or config['critic_lr'] is not None or config['weight_decay'] is not None or config['critic_weight_decay'] is not None:
+        if config['lr'] is not None or config['critic_lr'] is not None or config['weight_decay'] is not None or config['critic_weight_decay'] is not None or (config['separate_memory_lr'] and config['memory_lr'] is not None):
             self.optimizer, self.critic_optimizer = self._init_optimizers(
-                lr=config.get('lr', self.base_optim_config['lr']),
+                lr=(config.get('lr', self.base_optim_config['lr'][0]), config.get('memory_lr', self.base_optim_config['lr'][1])) if config.get('separate_memory_lr', False) else config.get('lr', self.base_optim_config['lr']),
                 critic_lr=config.get('critic_lr', self.base_optim_config['critic_lr']),
                 weight_decay=config.get('weight_decay', self.base_optim_config['weight_decay']),
-                critic_weight_decay=config.get('critic_weight_decay', self.base_optim_config['critic_weight_decay'])
+                critic_weight_decay=config.get('critic_weight_decay', self.base_optim_config['critic_weight_decay']),
+                separate_memory_lr=config.get('separate_memory_lr', False),
             )
 
         # 2. Get epochs and random resets configs

rxnn/transformers/layers.py

@@ -64,6 +64,13 @@ class ReactiveTransformerLayer(nn.Module):
             for param in self.norm2.parameters():
                 param.requires_grad_(is_trainable)
 
+    def memory_parameters(self) -> list[nn.Parameter]:
+        return list(self.memory_cross_attention.parameters()) + list(self.norm2.parameters())
+
+    def not_memory_parameters(self) -> list[nn.Parameter]:
+        memory_params = self.memory_parameters()
+        return [param for param in self.parameters() if param not in memory_params]
+
     def update_max_len(self, max_seq_len: int):
         if self.attention.rope is not None:
             self.attention.rope.update_max_len(max_seq_len)

rxnn/transformers/models.py

@@ -39,6 +39,16 @@ class ReactiveTransformerBase(nn.Module):
         for i in range(self.num_own_layers):
             self.layers[i].trainable_cross_attention_(is_trainable, with_norms)
 
+    def memory_parameters(self) -> list[nn.Parameter]:
+        own = [param for layer in self.layers for param in layer.memory_parameters()]
+        shared = [param for layer in self.shared_layers for param in layer.memory_parameters()] if self.shared_layers else []
+        return own + shared
+
+    def not_memory_parameters(self) -> list[nn.Parameter]:
+        own = [param for layer in self.layers for param in layer.not_memory_parameters()]
+        shared = [param for layer in self.shared_layers for param in layer.not_memory_parameters()] if self.shared_layers else []
+        return own + shared
+
     def moe_router_loss(self):
         return torch.stack([self.layers[i].moe_router_loss() for i in range(self.num_own_layers) if self.layers[i].use_moe or self.layers[i].use_moe_att] + [
             self.shared_layers[i].moe_router_loss() for i in range(self.num_shared_layers) if self.shared_layers[i].use_moe or self.shared_layers[i].use_moe_att]).mean()

{rxnn-0.2.24.dist-info → rxnn-0.2.26.dist-info}/METADATA

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

{rxnn-0.2.24.dist-info → rxnn-0.2.26.dist-info}/RECORD

@@ -5,18 +5,18 @@ rxnn/experimental/attention.py,sha256=46qwZLJuZMpIBrZ-r9DaQEPPmmZkO464C3Tkm_Mq-c
 rxnn/experimental/models.py,sha256=foBo0n0ufvBnfIdJomiEg3CuSOiWSt-q5ako7vzYxx4,4888
 rxnn/experimental/moe.py,sha256=jHZ1QhpWiVQOswVpFmuH7b2IUOPf0Uuf-I2Ddwsd7Us,6140
 rxnn/memory/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/memory/attention.py,sha256=p-r8DK3iVhNn-JAESVzIXDCG8gk1R_-x5xHclZ5jgb0,1813
-rxnn/memory/norm.py,sha256=Ofl8Q5NYEF9GQeO0bhM43tkTW91J0y6TSvTAOYMgloM,6278
+rxnn/memory/attention.py,sha256=PXVBZQYNsRraZh7QDBgUOdPy3lTI8B0d8CzduojBjG0,1747
+rxnn/memory/norm.py,sha256=E98jOQEuIOFFhlkvS8s4fFN-D4tLO6vaOqnObv1oVmA,6592
 rxnn/memory/stm.py,sha256=IH_3INw7FdI013t56ui3Zq9GPUq-k3HeZGjx6BerS4g,3888
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-rxnn/rxt/models.py,sha256=3gCYD_OXvQc8GaXQvRCSj1OcYOSHayWlpP5lsg9wMMk,12389
+rxnn/rxt/models.py,sha256=r8wZeeNTC2VAhiiNe4y7LrbnB4wjFu_cupKiGkpdgjI,13002
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/training/base.py,sha256=_xik1GXE4RJ_nxwqLQ1ccXA5pRtBCi-jL-jeRFBdHBU,11851
 rxnn/training/bml.py,sha256=FJszaQXOLx2ZHBa1CQpyMrG8i4Kj14E-gzDAEK_Ei5k,17272
 rxnn/training/callbacks.py,sha256=-N0MQPpZQaUWCINdTOsjul4bDGbGr2JgQBqOIXBLS6o,35053
 rxnn/training/dataset.py,sha256=7hTilFWPpqUEc6zNcMqBPjxFKxCfvTKKF3E8tVlwccQ,51250
-rxnn/training/models.py,sha256=5fl1hESVj2Hakqz5to8ZJzw5Q4_RKZAUq2bn6nRiPV8,6045
-rxnn/training/mrl.py,sha256=14wx3pVha15B7eRWPRgoxRtV5dPtBI0yadIHOYZjX6k,43275
+rxnn/training/models.py,sha256=_TrFwrQ_m6NDPalrafd8faPRyCnDFFFtN_gfzavaCFs,6474
+rxnn/training/mrl.py,sha256=hDsKQTaQcEVmnJruD3TxHZJJzDWu5I6Rq2HVDLj8ADU,44747
 rxnn/training/reward.py,sha256=7MTVdNm5HnWmt6zFDi3TAYmnVSL_-24riOoY2F7z4x8,11290
 rxnn/training/rl.py,sha256=j-KNLoZjhaEKasYNOc8DxHtwvknAgAJFwvXKot6otFA,3272
 rxnn/training/scheduler.py,sha256=LcjU35mEwz2U5x3U6tLfeeYlBqMxbFSxYzJYuXkWbSY,1408
@@ -25,14 +25,14 @@ rxnn/training/utils.py,sha256=Bw8nZLKIt7NQpUVCYkb_79kWKChVFOYgYXwODo4SvNc,5718
 rxnn/transformers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/transformers/attention.py,sha256=d0Igo1Nrn76BphbHrzekiKJfT3RCy4iSXSB6FLAOTgc,16247
 rxnn/transformers/ff.py,sha256=jJnuBDsnnX5uYC_WZH8cXAYrMnz0P-iX7MwcPivjRtI,2533
-rxnn/transformers/layers.py,sha256=UQZbrAg1UAttPASeqS7BP1a4JalktThmRMzX99Qghss,7618
+rxnn/transformers/layers.py,sha256=LXSY829fIHSCmFmClhQ6B7I5aKbiOqy9mZmwlJG_r7U,7961
 rxnn/transformers/mask.py,sha256=J0cfLVLt3SzS2ra3KcY4khrkhI975Dw4CjpUi3Sn25s,419
-rxnn/transformers/models.py,sha256=_2qO1SASHtKvTW3dW-Dy9HEmAvoNVC1_addm2tM9Zbs,8325
+rxnn/transformers/models.py,sha256=QwVxYN9DrKllEpOiFoAx4CiThOWafeTa-OAY7L6gN0Y,8929
 rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=t6iiQTdLQ0TakUWnnhKkb5DKF2F_9-thXHBydDF3fxg,17389
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.24.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.24.dist-info/METADATA,sha256=PrVfcCd8NBFtFnD8lAJqU7UW3lLEc-Tr7MQhK6obvuo,25960
-rxnn-0.2.24.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.24.dist-info/RECORD,,
+rxnn-0.2.26.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.26.dist-info/METADATA,sha256=XDqI42X3zLRAAKZlVLmstm24KFPP_MfvDtObG9GBc0Y,25960
+rxnn-0.2.26.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.26.dist-info/RECORD,,