rxnn 0.2.24__py3-none-any.whl → 0.2.25__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

@@ -20,8 +20,8 @@ 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
+        self.scale = nn.Parameter(torch.ones(num_slots, dim)) 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))
@@ -31,28 +31,26 @@ class AdaptivePositionalMemoryNorm(nn.Module):
             nn.init.normal_(self.scale, mean=1.0, std=0.01)
 
     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, 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 +75,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:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: rxnn
-Version: 0.2.24
+Version: 0.2.25
 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.25.dist-info}/RECORD

@@ -5,8 +5,8 @@ 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=mu_6iZJe61ag627csfJN2JK6QmmzofjOEhxV4ZWblXs,6410
 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
@@ -32,7 +32,7 @@ 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.25.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.25.dist-info/METADATA,sha256=nuGFk4oqSMhn6vrw2KZs4RtY0_ZLowg29IlkNVHZ6Jo,25960
+rxnn-0.2.25.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.25.dist-info/RECORD,,