rxnn 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

rxnn/memory/stm.py

@@ -5,7 +5,7 @@ class ShortTermMemory(nn.Module):
     """Short-term memory module for the Attention-based Memory System"""
 
     def __init__(self, num_layers: int, embed_dim: int, stm_size: int, init_type: str = 'normal',
-                 is_trainable: bool = False, legacy_init: bool = True, *args, **kwargs):
+                 is_trainable: bool = False, *args, **kwargs):
         super(ShortTermMemory, self).__init__(*args, **kwargs)
         self.num_layers = num_layers
         self.embed_dim = embed_dim
@@ -14,10 +14,6 @@ class ShortTermMemory(nn.Module):
         self.is_trainable = is_trainable
         assert init_type in ['normal', 'standard', 'uniform', 'ones', 'zeros'], \
             'STM init type must be one of "normal", "standard", "uniform", "ones", "zeros"'
-
-        # Legacy init - temporary option to load old models with not-batched STM (they will be loaded, updated and then the option will be removed)
-        self.legacy_init = legacy_init
-
         self.init_type = init_type
         stm = self._init_tensor()
         if self.is_trainable:
@@ -27,8 +23,7 @@ class ShortTermMemory(nn.Module):
 
     def _init_tensor(self, init_type: str = None):
         init_type = init_type or self.init_type
-        stm_shape = (self.num_layers, self.stm_size, self.embed_dim) \
-            if self.legacy_init else (self.num_layers, self.batch_size, self.stm_size, self.embed_dim)
+        stm_shape = (self.num_layers, self.batch_size, self.stm_size, self.embed_dim)
         if init_type == 'normal':
             return torch.normal(0, 0.02, stm_shape)
         elif init_type == 'standard':
@@ -40,12 +35,8 @@ class ShortTermMemory(nn.Module):
         else:
             return torch.zeros(*stm_shape)
 
-    def reset_legacy_(self):
-        self.legacy_init = False
-        self.memory = self._init_tensor()
-
     def forward(self, layer: int) -> torch.Tensor:
-        return self.memory[layer].unsqueeze(0) if self.legacy_init else self.memory[layer]
+        return self.memory[layer]
 
     def update_layer(self, layer: int, new_stm: torch.Tensor):
         self.memory[layer] = new_stm

{rxnn-0.2.1.dist-info → rxnn-0.2.2.dist-info}/METADATA

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

{rxnn-0.2.1.dist-info → rxnn-0.2.2.dist-info}/RECORD

@@ -7,7 +7,7 @@ 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/stm.py,sha256=tEAC5c4QPl-ykNm-k2DlcP6Q2NQKGXMnpKLCx92lI30,4083
+rxnn/memory/stm.py,sha256=eSMK5KdupWNf56FcDYprHnjA51EeYBzSKza7tiZxKSc,3618
 rxnn/rxt/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rxnn/rxt/models.py,sha256=zNrf6mn-s2vJyauHwNgYm_e-gFI1clmXp_JyCKGQD3E,12083
 rxnn/training/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -32,7 +32,7 @@ rxnn/transformers/moe.py,sha256=j6jEx6Ip0zttlUZKKn82azxo95lkLZs-H2GLSMD88hY,5859
 rxnn/transformers/positional.py,sha256=1PjcJybUzeQlIKJI4tahAGZcYgCRCL0otxs7mpsNuzM,4410
 rxnn/transformers/sampler.py,sha256=2dpUQv88ekZa_CMSPLrXvB6X684wxUE2bDVznsi5ACs,17429
 rxnn/utils.py,sha256=ihb6OTyDtPiocB_lOvnq7eOkjjpCkgs8wxvXUBNQ7mM,996
-rxnn-0.2.1.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
-rxnn-0.2.1.dist-info/METADATA,sha256=0NaOW0HNG4nGR3Vchlp-vlzzx7NBSFpAPFYwGzT3qP8,25959
-rxnn-0.2.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-rxnn-0.2.1.dist-info/RECORD,,
+rxnn-0.2.2.dist-info/LICENSE,sha256=C8coDFIUYuOcke4JLPwTqahQUCyXyGq6WOaigOkx8tY,11275
+rxnn-0.2.2.dist-info/METADATA,sha256=GlH7tyaDt27dzlp7G3CafWLAic8S5dTd-eiYKzDNQlA,25959
+rxnn-0.2.2.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+rxnn-0.2.2.dist-info/RECORD,,