x-transformers 2.4.5__py3-none-any.whl → 2.4.7__py3-none-any.whl

x_transformers/up_wrapper.py

@@ -6,7 +6,7 @@ from functools import partial
 from random import randrange, uniform
 
 import torch
-from torch import nn, cat, randperm
+from torch import nn, cat, tensor, randperm
 from torch.nn import LSTM, Module
 
 from x_transformers.x_transformers import (
@@ -76,6 +76,13 @@ class SyntheticDataGenerator(Module):
 
         self.apply(self.init_)
 
+    def reset_(self):
+        for m in self.modules():
+            if hasattr(m, 'reset_parameters'):
+                m.reset_parameters()
+
+        self.apply(self.init_)
+
     @torch.no_grad()
     def init_(self, m):
         if isinstance(m, nn.Linear):
@@ -138,7 +145,9 @@ class UniversalPretrainWrapper(Module):
         num_reset = 20,
         batch_size = 32,
         seq_len = 512,
-        seed_length = 8
+        seed_length = 8,
+        reset_turing_machine_every = 0,
+        keep_buffer_on_cpu = False
     ):
         super().__init__()
 
@@ -157,6 +166,8 @@ class UniversalPretrainWrapper(Module):
                 max_seq_len = seq_len
             )
 
+        self.reset_turing_machine_every = reset_turing_machine_every
+
         self.seq_len = seq_len
         self.data_generator = data_generator
 
@@ -175,11 +186,16 @@ class UniversalPretrainWrapper(Module):
 
         init_data_buffer = self.random_sequences_fn(buffer_size // 2, buffer_size // 2)
 
-        self.register_buffer('synth_data_buffer', init_data_buffer)
+        if keep_buffer_on_cpu:
+            self.synth_data_buffer = init_data_buffer
+        else:
+            self.register_buffer('synth_data_buffer', init_data_buffer)
+
+        self.register_buffer('step', tensor(0))
 
     @property
     def device(self):
-        return self.synth_data_buffer.device
+        return self.step.device
 
     def get_rand_sequences_from_buffer(self, size = None):
         size = default(size, self.batch_size)
@@ -206,10 +222,17 @@ class UniversalPretrainWrapper(Module):
 
         generated = self.data_generator.generate(
             self.seq_len,
-            condition = conditions,
-            seed = seeds
+            condition = conditions.to(self.device),
+            seed = seeds.to(self.device)
         )
 
+        self.step.add_(1)
+
+        # maybe reset turing machine
+
+        if self.reset_turing_machine_every > 0 and divisible_by(self.step.item(), self.reset_turing_machine_every):
+            self.data_generator.reset_()
+
         # reset
 
         if self.num_reset > 0:
@@ -226,6 +249,6 @@ class UniversalPretrainWrapper(Module):
 
         # sample yet again according to pseudocode
 
-        data = self.get_rand_sequences_from_buffer()
+        data = self.get_rand_sequences_from_buffer().to(self.device)
 
         return self.ar_wrapped(data)

{x_transformers-2.4.5.dist-info → x_transformers-2.4.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.4.5
+Version: 2.4.7
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers

{x_transformers-2.4.5.dist-info → x_transformers-2.4.7.dist-info}/RECORD

@@ -8,11 +8,11 @@ x_transformers/entropy_based_tokenizer.py,sha256=F2lO8-v3aLIcVDVNhu7RR-UtRdlmaaY
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=2NU58hYMgn-4Jzg3mie-mXb0XH_dCN7fjlzd3K1rLUY,10510
-x_transformers/up_wrapper.py,sha256=RW-0EmheBqJz7g4j-M37XxWMttz-QuPhmnRWMrD7_-0,6103
+x_transformers/up_wrapper.py,sha256=UDaLCQSeoFtC2b9ZPDbof4pzpbEBPO-FjQuLiA2kjP0,6839
 x_transformers/x_transformers.py,sha256=9Fi0HvzpeIJqM6HlAd2M6JqsfjhTN1zEH9iFIimyjS4,117608
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.4.5.dist-info/METADATA,sha256=J7sgvMp5pLgUWZ37qdDF4EBmECBqfgEZ0rSC7w12AAI,90223
-x_transformers-2.4.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.4.5.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.4.5.dist-info/RECORD,,
+x_transformers-2.4.7.dist-info/METADATA,sha256=W-Er0p8u3zxuJmYnA8qA6u1GDUqB8p2yTlUS6Xhqsb0,90223
+x_transformers-2.4.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.4.7.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.4.7.dist-info/RECORD,,