x-transformers 1.26.3__py3-none-any.whl → 1.26.4__py3-none-any.whl

x_transformers/autoregressive_wrapper.py

@@ -145,7 +145,7 @@ class AutoregressiveWrapper(Module):
         cache_kv = True,
         **kwargs
     ):
-        max_seq_len, device = self.max_seq_len, prompts.device
+        max_seq_len, greedy, device = self.max_seq_len, temperature == 0., prompts.device
 
         prompts, ps = pack([prompts], '* n')
 
@@ -230,11 +230,14 @@ class AutoregressiveWrapper(Module):
 
             # filter by top_k, top_p (nucleus), top_a, or custom
 
-            filtered_logits = filter_logits_fn(logits, **filter_kwargs)
+            if greedy:
+                sample = logits.argmax(dim = -1, keepdim = True)
+            else:
+                filtered_logits = filter_logits_fn(logits, **filter_kwargs)
+                probs = F.softmax(filtered_logits / temperature, dim=-1)
+                sample = torch.multinomial(probs, 1)
 
-            probs = F.softmax(filtered_logits / temperature, dim=-1)
-
-            sample = torch.multinomial(probs, 1)
+            # concat sample
 
             out = torch.cat((out, sample), dim=-1)
 

{x_transformers-1.26.3.dist-info → x_transformers-1.26.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: x-transformers
-Version: 1.26.3
+Version: 1.26.4
 Summary: X-Transformers - Pytorch
 Home-page: https://github.com/lucidrains/x-transformers
 Author: Phil Wang

{x_transformers-1.26.3.dist-info → x_transformers-1.26.4.dist-info}/RECORD

@@ -1,13 +1,13 @@
 x_transformers/__init__.py,sha256=pXc_U4M3ONUQcpNgZySDIlCF1rp7u4FFmcOYjc4WuXw,629
 x_transformers/attend.py,sha256=MFl_FbgPsm9mziZPTi_s8QbxASETwbGeciMH8sUIwT8,10188
-x_transformers/autoregressive_wrapper.py,sha256=6O4fz0keP2EBaRssUda7I5ZFxY-LFCzM-m--NtQp9rw,9058
+x_transformers/autoregressive_wrapper.py,sha256=xqr-CnB98N89y7K0z2SLEy8If8pE_3MP7C9WFTwjefs,9230
 x_transformers/continuous.py,sha256=ixfgi2_zpGN03SX_STXFkNYEOAkgwVIxuS53QgDCx-g,6026
 x_transformers/nonautoregressive_wrapper.py,sha256=AQLE4rA_Kh8VNoe9OzpwyeWson34sRkhks4dn4seNjI,10414
 x_transformers/x_transformers.py,sha256=gXq_IpWswCeM06r0VCxvZGTVndsxD69_uvKSJFC1544,61632
 x_transformers/xl_autoregressive_wrapper.py,sha256=DCx4n0_c1tFai4nOqaWVnqx2p9eutsZsDMiMP1ckxNU,4117
 x_transformers/xval.py,sha256=lS9W_E_RskPQAqVZkPiUzbByoW1Ajsw_phsikA3JDAg,8139
-x_transformers-1.26.3.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-1.26.3.dist-info/METADATA,sha256=K6TZ-TdM0z8jBSMvLMJjBI7nVTtp7rKHW-5wddYYHYQ,661
-x_transformers-1.26.3.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-x_transformers-1.26.3.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
-x_transformers-1.26.3.dist-info/RECORD,,
+x_transformers-1.26.4.dist-info/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-1.26.4.dist-info/METADATA,sha256=fX_dEQcxjW00biyG6pOz9J88HGiS9qVS05KP1Z94H44,661
+x_transformers-1.26.4.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+x_transformers-1.26.4.dist-info/top_level.txt,sha256=hO6KGpFuGucRNEtRfme4A_rGcM53AKwGP7RVlRIxS5Q,15
+x_transformers-1.26.4.dist-info/RECORD,,