x-transformers 2.10.1__tar.gz → 2.11.0__tar.gz

{x_transformers-2.10.1 → x_transformers-2.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.10.1
+Version: 2.11.0
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -2598,4 +2598,13 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@inproceedings{Fleuret2025TheFT,
+    title   = {The Free Transformer},
+    author  = {Franccois Fleuret},
+    year    = {2025},
+    url     = {https://api.semanticscholar.org/CorpusID:282210283}
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.10.1 → x_transformers-2.11.0}/README.md

@@ -2549,4 +2549,13 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@inproceedings{Fleuret2025TheFT,
+    title   = {The Free Transformer},
+    author  = {Franccois Fleuret},
+    year    = {2025},
+    url     = {https://api.semanticscholar.org/CorpusID:282210283}
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.10.1 → x_transformers-2.11.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "x-transformers"
-version = "2.10.1"
+version = "2.11.0"
 description = "X-Transformers"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{x_transformers-2.10.1 → x_transformers-2.11.0}/tests/test_x_transformers.py

@@ -1,4 +1,5 @@
 import pytest
+param = pytest.mark.parametrize
 
 import torch
 from torch import nn
@@ -186,7 +187,7 @@ def test_average_pool_embed():
 
     assert logits.shape == (2, 20000)
 
-@pytest.mark.parametrize('num_cls_tokens', (1, 2))
+@param('num_cls_tokens', (1, 2))
 def test_cls_token(num_cls_tokens):
     model = TransformerWrapper(
         num_tokens = 20000,
@@ -234,7 +235,7 @@ def test_squeeze_logit_dim_one():
 
     assert logits.shape == (2,)
 
-@pytest.mark.parametrize('depth', (4, 5))
+@param('depth', (4, 5))
 def test_unet_skip(depth):
 
     model = TransformerWrapper(
@@ -294,7 +295,7 @@ def test_mos():
 
     eval_logits = model(x)
 
-@pytest.mark.parametrize('attn_one_kv_head', (True, False))
+@param('attn_one_kv_head', (True, False))
 def test_l2_distance(attn_one_kv_head):
 
     model = TransformerWrapper(
@@ -331,7 +332,7 @@ def test_reinject_input():
 
     model(x) # (1, 1024, 20000)
 
-@pytest.mark.parametrize('learned_value_residual_mix', (False, True))
+@param('learned_value_residual_mix', (False, True))
 def test_value_residual(
     learned_value_residual_mix: bool
 ):
@@ -352,7 +353,7 @@ def test_value_residual(
 
     model(x)
 
-@pytest.mark.parametrize('has_num_mem_kv', (False, True))
+@param('has_num_mem_kv', (False, True))
 def test_forgetting_transformer(
     has_num_mem_kv: bool
 ):
@@ -388,7 +389,7 @@ def test_neo_mlp():
     out = mlp(x)
     assert out.shape == (3, 7)
 
-@pytest.mark.parametrize('flash', (True, False))
+@param('flash', (True, False))
 def test_custom_alibi(flash: bool):
 
     model = TransformerWrapper(
@@ -409,7 +410,7 @@ def test_custom_alibi(flash: bool):
 
     logits = model(x, pos = pos)
 
-@pytest.mark.parametrize('rotary_xpos', (True, False))
+@param('rotary_xpos', (True, False))
 def test_custom_rotary_pos_emb(rotary_xpos):
     from einops import repeat
 
@@ -433,7 +434,7 @@ def test_custom_rotary_pos_emb(rotary_xpos):
     logits2 = model(x)
     assert torch.allclose(logits1, logits2)
 
-@pytest.mark.parametrize('flash', (True, False))
+@param('flash', (True, False))
 def test_custom_alibi_across_heads(flash: bool):
     model = Decoder(
         dim = 512,
@@ -455,7 +456,7 @@ def test_custom_alibi_across_heads(flash: bool):
 
     embed = model(x, pos = pos)
 
-@pytest.mark.parametrize('embedder_type', ('embedding', 'none', 'custom'))
+@param('embedder_type', ('embedding', 'none', 'custom'))
 def test_embedder(embedder_type):
     num_tokens = 20000
     dim = 128
@@ -502,7 +503,7 @@ def test_embedder(embedder_type):
     assert output.shape == (2, 1024, 20000)
 
 
-@pytest.mark.parametrize("to_logits", ('linear', 'none', 'pointer'))
+@param("to_logits", ('linear', 'none', 'pointer'))
 def test_to_logits(to_logits):
     num_tokens = 20000
     dim = 128
@@ -560,8 +561,8 @@ def test_laser():
 
     model(x)
 
-@pytest.mark.parametrize('self_attn_custom_pos', (True, False))
-@pytest.mark.parametrize('cross_attn_rotary', (True, False))
+@param('self_attn_custom_pos', (True, False))
+@param('cross_attn_rotary', (True, False))
 def test_cross_attn_rotary(
     self_attn_custom_pos: bool,
     cross_attn_rotary: bool
@@ -593,7 +594,7 @@ def test_cross_attn_rotary(
       context_mask = context_mask
     )
 
-@pytest.mark.parametrize('tanh', (True, False))
+@param('tanh', (True, False))
 def test_hyper_connections(tanh):
 
     model = TransformerWrapper(
@@ -614,7 +615,7 @@ def test_hyper_connections(tanh):
 
     model(x)
 
-@pytest.mark.parametrize('hybrid_axial_dim', (1, 4))
+@param('hybrid_axial_dim', (1, 4))
 def test_hybrid(hybrid_axial_dim):
     from torch.nn import GRU
 
@@ -770,8 +771,8 @@ def test_multi_latent_attention():
 
     model(x)
 
-@pytest.mark.parametrize('num_residual_streams', (1, 4))
-@pytest.mark.parametrize('integrate_layers', (False, True))
+@param('num_residual_streams', (1, 4))
+@param('integrate_layers', (False, True))
 def test_lime(
     num_residual_streams,
     integrate_layers
@@ -792,10 +793,10 @@ def test_lime(
 
     model(x)
 
-@pytest.mark.parametrize('backward_ar_loss_weight', (1., 0.5))
-@pytest.mark.parametrize('goal_suffix', (False, True))
-@pytest.mark.parametrize('pred_distance', (False, True))
-@pytest.mark.parametrize('variable_len', (False, True))
+@param('backward_ar_loss_weight', (1., 0.5))
+@param('goal_suffix', (False, True))
+@param('pred_distance', (False, True))
+@param('variable_len', (False, True))
 def test_belief_state_wrapper(
     backward_ar_loss_weight,
     goal_suffix,
@@ -867,7 +868,7 @@ def test_dynamic_tanh():
 
     model(x)
 
-@pytest.mark.parametrize('var_length', (False, True))
+@param('var_length', (False, True))
 def test_entropy_based_tokenizer(
     var_length
 ):
@@ -966,9 +967,9 @@ def test_ff_deep_embed():
 
     assert logits.shape == (2, 1024, 20000)
 
-@pytest.mark.parametrize('probabilistic', (False, True))
-@pytest.mark.parametrize('cache_kv', (False, True))
-@pytest.mark.parametrize('rollout_steps', (1, 4))
+@param('probabilistic', (False, True))
+@param('cache_kv', (False, True))
+@param('rollout_steps', (1, 4))
 def test_continuous(
     probabilistic,
     cache_kv,
@@ -1012,7 +1013,7 @@ def test_continuous(
     generated = model.generate(start_emb, 17, cache_kv = cache_kv) # (17, 777)
     assert generated.shape == (17, 777)
 
-@pytest.mark.parametrize('add_continuous_pred_head', (False, True))
+@param('add_continuous_pred_head', (False, True))
 def test_autoregressive_wrapper(
     add_continuous_pred_head
 ):
@@ -1100,7 +1101,7 @@ def add_attn_pool():
 
     assert intermediates.attn_pooled_tokens.shape[1] == 3
 
-@pytest.mark.parametrize('keep_buffer_on_cpu', (False, True))
+@param('keep_buffer_on_cpu', (False, True))
 def test_up(
     keep_buffer_on_cpu
 ):
@@ -1126,7 +1127,7 @@ def test_up(
     loss = up_wrapper()
     loss.backward()
 
-@pytest.mark.parametrize('stochastic', (False, True))
+@param('stochastic', (False, True))
 def test_beam_search(stochastic):
     from x_transformers import TransformerWrapper, Decoder, AutoregressiveWrapper
 
@@ -1154,8 +1155,8 @@ def test_beam_search(stochastic):
     assert scores.shape == (4, 2)
 
 
-@pytest.mark.parametrize('num_pooled_tokens', (1, 3))
-@pytest.mark.parametrize('attn_pool_depth', (1, 3))
+@param('num_pooled_tokens', (1, 3))
+@param('attn_pool_depth', (1, 3))
 def test_attn_pooler(
     num_pooled_tokens,
     attn_pool_depth
@@ -1288,7 +1289,7 @@ def test_accept_layer_intermediates():
 
     assert embeds.shape == (3, 32, 512)
 
-@pytest.mark.parametrize('use_loss_weight', (False, True))
+@param('use_loss_weight', (False, True))
 def test_simple_mdlm(
     use_loss_weight
 ):
@@ -1386,7 +1387,10 @@ def test_stochastic_attn():
     log_probs = log_prob_from_hard_attend(intermediate)
     assert log_probs.shape == (1, 8, 1024)
 
-def test_attn_negative_weights():
+@param('head_learned_sink', (True, False))
+def test_attn_negative_weights(
+    head_learned_sink
+):
     from x_transformers import TransformerWrapper, Decoder
 
     model = TransformerWrapper(
@@ -1396,7 +1400,8 @@ def test_attn_negative_weights():
             dim = 512,
             depth = 12,
             heads = 8,
-            attn_cog_signed = True
+            attn_cog_signed = True,
+            attn_head_learned_sink = True
         ),
     )
 

{x_transformers-2.10.1 → x_transformers-2.11.0}/train_copy.py

@@ -12,32 +12,35 @@ GENERATE_EVERY  = 100
 NUM_TOKENS = 16 + 2
 ENC_SEQ_LEN = 32
 DEC_SEQ_LEN = 64 + 1
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 
 # helpers
 
 def cycle():
     while True:
-        prefix = torch.ones((BATCH_SIZE, 1)).long().cuda()
-        src = torch.randint(2, NUM_TOKENS, (BATCH_SIZE, ENC_SEQ_LEN)).long().cuda()
+        prefix = torch.ones((BATCH_SIZE, 1)).long().to(DEVICE)
+        src = torch.randint(2, NUM_TOKENS, (BATCH_SIZE, ENC_SEQ_LEN)).long().to(DEVICE)
         tgt = torch.cat((prefix, src, src), 1)
-        src_mask = torch.ones(BATCH_SIZE, src.shape[1]).bool().cuda()
+        src_mask = torch.ones(BATCH_SIZE, src.shape[1]).bool().to(DEVICE)
         yield (src, tgt, src_mask)
 
 # instantiate model
 
 model = XTransformer(
-    dim = 512,
+    dim = 128,
     tie_token_emb = True,
     return_tgt_loss = True,
     enc_num_tokens=NUM_TOKENS,
     enc_depth = 3,
     enc_heads = 8,
     enc_max_seq_len = ENC_SEQ_LEN,
+    enc_attn_cog_signed = True,
     dec_num_tokens = NUM_TOKENS,
     dec_depth = 3,
     dec_heads = 8,
-    dec_max_seq_len = DEC_SEQ_LEN
-).cuda()
+    dec_max_seq_len = DEC_SEQ_LEN,
+    dec_attn_cog_signed = True
+).to(DEVICE)
 
 # optimizer
 
@@ -61,10 +64,10 @@ for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10., desc='training'):
         model.eval()
         src, _, src_mask = next(cycle())
         src, src_mask = src[:1], src_mask[:1]
-        start_tokens = (torch.ones((1, 1)) * 1).long().cuda()
+        start_tokens = (torch.ones((1, 1)) * 1).long().to(DEVICE)
 
         sample = model.generate(src, start_tokens, ENC_SEQ_LEN, mask = src_mask)
-        incorrects = (src != sample).abs().sum()
+        incorrects = (src != sample).long().abs().sum()
 
         print(f"input:  ", src)
         print(f"predicted output:  ", sample)

x_transformers-2.11.0/train_free.py

@@ -0,0 +1,134 @@
+# /// script
+# dependencies = [
+#   "tqdm",
+#   "x-transformers>=2.11.0",
+# ]
+# ///
+
+from x_transformers.free_transformer import FreeTransformer
+
+from math import log
+import random
+import tqdm
+import gzip
+import numpy as np
+
+import torch
+import torch.optim as optim
+from torch import tensor
+from torch.nn import functional as F
+from torch.utils.data import DataLoader, Dataset
+
+# constants
+
+NUM_BATCHES = int(1e5)
+BATCH_SIZE = 4
+GRADIENT_ACCUMULATE_EVERY = 4
+LEARNING_RATE = 1e-4
+VALIDATE_EVERY  = 100
+GENERATE_EVERY  = 250
+GENERATE_LENGTH = 512
+PRIME_LENGTH = 32
+SEQ_LEN = 512
+
+LATENT_BITS = 8
+NAT = log(2)
+
+# helpers
+
+def cycle(loader):
+    while True:
+        for data in loader:
+            yield data
+
+def decode_token(token):
+    return str(chr(max(32, token)))
+
+def decode_tokens(tokens):
+    return ''.join(list(map(decode_token, tokens)))
+
+# instantiate GPT-like decoder model
+
+model = FreeTransformer(
+    num_tokens = 256,
+    max_seq_len = SEQ_LEN,
+    dim = 512,
+    heads = 8,
+    rotary_pos_emb = True,
+    dec_head_depth = 4,
+    dec_tail_depth = 4,
+    enc_depth = 3,
+    kl_loss_weight = 1.,
+    kl_loss_threshold = NAT,
+    latent_bits = LATENT_BITS
+).cuda()
+
+rand_index = torch.randint(0, 2 ** LATENT_BITS, ())
+latents = F.one_hot(rand_index, 2 ** LATENT_BITS).float().cuda()
+
+# prepare enwik8 data
+
+with gzip.open('./data/enwik8.gz') as file:
+    data = np.frombuffer(file.read(int(95e6)), dtype=np.uint8).copy()
+    train_x, valid_x = np.split(data, [int(90e6)])
+    data_train, data_val = torch.from_numpy(train_x), torch.from_numpy(valid_x)
+
+class TextSamplerDataset(Dataset):
+    def __init__(self, data, seq_len):
+        super().__init__()
+        self.data = data
+        self.seq_len = seq_len
+
+    def __getitem__(self, index):
+        rand_start = torch.randint(0, self.data.size(0) - self.seq_len - 1, (1,))
+        full_seq = self.data[rand_start: rand_start + self.seq_len + 1].long()
+        return full_seq.cuda()
+
+    def __len__(self):
+        return self.data.size(0) // self.seq_len
+
+train_dataset = TextSamplerDataset(data_train, SEQ_LEN)
+val_dataset   = TextSamplerDataset(data_val, SEQ_LEN)
+train_loader  = cycle(DataLoader(train_dataset, batch_size = BATCH_SIZE, drop_last = True))
+val_loader    = cycle(DataLoader(val_dataset, batch_size = BATCH_SIZE, drop_last = True))
+
+# optimizer
+
+optim = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)
+
+# training
+
+for i in tqdm.tqdm(range(NUM_BATCHES), mininterval=10., desc='training'):
+    model.train()
+
+    for __ in range(GRADIENT_ACCUMULATE_EVERY):
+        loss, (ar_loss, vae_kl_loss) = model(next(train_loader), return_all_losses = True)
+        (loss / GRADIENT_ACCUMULATE_EVERY).backward()
+
+    print(f'training loss: {ar_loss.item():.4f}\t| kl loss: {vae_kl_loss.item():.4f}')
+
+    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+    optim.step()
+    optim.zero_grad()
+
+    if i % VALIDATE_EVERY == 0:
+        model.eval()
+        with torch.no_grad():
+            loss, (ar_loss, _) = model(next(val_loader), return_all_losses = True)
+            print(f'validation loss: {ar_loss.item():.4f}')
+
+    if i % GENERATE_EVERY == 0:
+        model.eval()
+        inp = random.choice(val_dataset)[:PRIME_LENGTH]
+        prime = decode_tokens(inp)
+        print(f'%s \n\n %s', (prime, '*' * 100))
+
+        sample = model.generate(
+            prompts = inp,
+            seq_len = GENERATE_LENGTH,
+            latents = latents
+        )
+
+        output_str = decode_tokens(sample)
+
+        print(f'\n\nlatent {rand_index.tolist()} - ', output_str)

{x_transformers-2.10.1 → x_transformers-2.11.0}/x_transformers/attend.py

@@ -549,6 +549,11 @@ class Attend(Module):
         if self.head_learned_sink:
             # add learned attention sink
             attn_sink = repeat(self.head_attn_sink, 'h -> b h i 1', b = sim.shape[0], i = sim.shape[2])
+
+            if self.cog_signed:
+                attn_sink, attn_sink_sign = attn_sink.abs(), attn_sink.sign()
+                sim_sign = cat((attn_sink_sign, sim_sign), dim = -1)
+
             sim = cat((attn_sink, sim), dim = -1)
 
         pre_softmax_attn = sim

{x_transformers-2.10.1 → x_transformers-2.11.0}/x_transformers/autoregressive_wrapper.py

@@ -43,6 +43,10 @@ def log(t, eps = 1e-20):
 def gumbel_noise(t):
     return -log(-log(torch.rand_like(t)))
 
+def gumbel_sample(logits, temperature = 1., eps = 1e-6):
+    noise = gumbel_noise(logits)
+    return ((logits / max(temperature, eps)) + noise).argmax(dim = -1)
+
 # function for modifying all the cached key / values
 
 def modify_cached_kv(cache, fn):

x_transformers-2.11.0/x_transformers/free_transformer.py

@@ -0,0 +1,330 @@
+from __future__ import annotations
+
+# https://arxiv.org/abs/2510.17558
+# François Fleuret
+# https://www.youtube.com/watch?v=Nao16-6l6dQ
+
+import math
+
+import torch
+from torch import nn, Tensor, is_tensor, tensor, arange
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList
+
+from x_transformers.x_transformers import (
+    Encoder,
+    Decoder,
+    TransformerWrapper
+)
+
+from x_transformers.autoregressive_wrapper import (
+    gumbel_sample,
+    top_p,
+    top_k
+)
+
+from einops.layers.torch import Rearrange, Reduce
+from einops import rearrange, reduce, repeat, einsum, pack, unpack
+
+# helper functions
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def log(t, eps = 1e-20):
+    return t.clamp_min(eps).log()
+
+def pack_with_inverse(t, pattern):
+    packed, ps = pack([t], pattern)
+
+    def inverse(out, inv_pattern = None):
+        inv_pattern = default(inv_pattern, pattern)
+        unpacked, = unpack(out, ps, inv_pattern)
+        return unpacked
+
+    return packed, inverse
+
+# binary mapper
+
+NAT = math.log(2)
+
+def binary_entropy(logits):
+    prob = logits.sigmoid()
+    not_prob = 1. - prob
+    return -(prob * F.logsigmoid(logits) + not_prob * F.logsigmoid(-logits)).sum(dim = -1)
+
+class BinaryMapper(Module):
+    def __init__(
+        self,
+        bits = 1,
+        kl_loss_threshold = NAT # 1 bit
+    ):
+        super().__init__()
+
+        self.bits = bits
+        self.num_codes = 2 ** bits
+        self.kl_loss_threshold = kl_loss_threshold
+
+        power_two = 2 ** arange(bits)
+        codes = (arange(self.num_codes)[:, None].bitwise_and(power_two) != 0).byte().bool()
+
+        self.register_buffer('power_two', power_two, persistent = False)
+        self.register_buffer('codes', codes, persistent = False)
+
+    def forward(
+        self,
+        logits,
+        temperature = 1.,
+        straight_through = None
+    ):
+        straight_through = default(straight_through, self.training)
+
+        assert logits.shape[-1] == self.bits, f'logits must have a last dimension of {self.bits}'
+
+        # temperature and prob for sampling
+
+        prob_for_sample = (logits / temperature).sigmoid()
+
+        # sampling
+
+        sampled_bits = (torch.rand_like(logits) <= prob_for_sample).long()
+        indices = (self.power_two * sampled_bits).sum(dim = -1)
+
+        one_hot = F.one_hot(indices, self.num_codes).float()
+
+        # return hard one hot if not training or overridden
+
+        if not straight_through:
+            return one_hot
+
+        # calculate negative entropy
+
+        kl_div = self.bits * NAT - binary_entropy(logits)
+        aux_kl_loss = F.relu(kl_div - self.kl_loss_threshold).mean()
+
+        # get the soft G for the gradients and do a straight through
+
+        soft_G = (
+            einsum(F.logsigmoid(logits), self.codes.float(), '... bits, codes bits -> ... codes') +
+            einsum(F.logsigmoid(-logits), (~self.codes).float(), '... bits, codes bits -> ... codes')
+        ).exp()
+
+        # straight through
+
+        one_hot = one_hot + soft_G - soft_G.detach()
+
+        return one_hot, aux_kl_loss
+
+# classes
+
+class FreeTransformer(Module):
+    def __init__(
+        self,
+        *,
+        num_tokens,
+        dim,
+        dec_head_depth,
+        dec_tail_depth,
+        enc_depth,
+        max_seq_len,
+        dim_latent = None,
+        attn_dim_head = 64,
+        heads = 8,
+        latent_bits = 16,
+        kl_loss_threshold = NAT,
+        binary_mapper_kwargs: dict = dict(),
+        enc_kwargs: dict = dict(),
+        dec_kwargs: dict = dict(),
+        kl_loss_weight = 1.,
+        pad_id = -1,
+        encoder: Module | None = None,
+        **kwargs
+    ):
+        super().__init__()
+        dim_latent = default(dim_latent, dim)
+
+        self.token_emb = nn.Embedding(num_tokens, dim)
+
+        self.token_unembed = nn.Linear(dim, num_tokens, bias = False)
+
+        if not exists(encoder):
+            encoder = Encoder(
+                dim = dim,
+                depth = enc_depth,
+                attn_dim_head = attn_dim_head,
+                heads = heads,
+                **kwargs,
+                **enc_kwargs
+            )
+
+        self.encoder = encoder
+
+        self.to_latent_bit_logits = nn.Sequential(
+            Reduce('b n d -> b d', 'mean'),
+            nn.Linear(dim, latent_bits, bias = False),
+        )
+
+        self.binary_mapper = BinaryMapper(
+            latent_bits,
+            kl_loss_threshold,
+            **binary_mapper_kwargs
+        )
+
+        self.from_latent_to_condition = nn.Sequential(
+            nn.Linear(2 ** latent_bits, dim, bias = False),
+            Rearrange('b d -> b 1 d')
+        )
+
+        self.decoder_head = Decoder(
+            dim = dim,
+            depth = dec_head_depth,
+            attn_dim_head = attn_dim_head,
+            heads = heads,
+            pre_norm_has_final_norm = False,
+            **kwargs,
+            **dec_kwargs
+        )
+
+        self.decoder_tail = Decoder(
+            dim = dim,
+            depth = dec_tail_depth,
+            attn_dim_head = attn_dim_head,
+            heads = heads,
+            pre_norm_has_final_norm = True,
+            **kwargs,
+            **dec_kwargs
+        )
+
+        self.pad_id = pad_id
+
+        self.kl_loss_weight = kl_loss_weight
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def encode_to_latents(
+        self,
+        seq,
+        mask = None,
+        return_kl_loss = False
+    ):
+        pooled = self.encoder(seq, mask = mask)
+
+        bit_logits = self.to_latent_bit_logits(pooled)
+
+        one_hot_latents, kl_loss = self.binary_mapper(bit_logits, straight_through = True)
+
+        if not return_kl_loss:
+            return one_hot_latents
+
+        return one_hot_latents, kl_loss
+
+    @torch.no_grad()
+    def generate(
+        self,
+        prompts,
+        seq_len,
+        latents = None,
+        filter_logits_fn = top_p,
+        logit_filter_kwargs: dict = dict(thres = 0.9)
+    ):
+        prompts, inverse_pack = pack_with_inverse(prompts, '* n')
+
+        batch = prompts.shape[0]
+
+        # prepend embeds
+
+        condition = None
+        if exists(latents):
+            if not is_tensor(latents):
+                latents = tensor(latents, device = self.device)
+
+            if latents.ndim == 1: # repeat latents
+                latents = repeat(latents, 'd -> b d', b = batch)
+
+            condition = self.from_latent_to_condition(latents)
+
+        # generated
+
+        prompt_len = prompts.shape[-1]
+
+        generated = prompts
+
+        tokens = self.token_emb(generated)
+
+        for _ in range(max(0, seq_len - prompt_len)):
+
+            head_embed = self.decoder_head(tokens)
+
+            if exists(condition):
+                head_embed = head_embed + condition
+
+            tail_embed = self.decoder_tail(head_embed)
+
+            tail_embed = tail_embed[:, -1]
+
+            logits = self.token_unembed(tail_embed)
+
+            logits = filter_logits_fn(logits, **logit_filter_kwargs)
+
+            sampled = gumbel_sample(logits)
+
+            generated, _ = pack((generated, sampled), 'b *')
+            tokens, _ = pack((tokens, self.token_emb(sampled)), 'b * d')
+
+        return inverse_pack(generated)
+
+    def forward(
+        self,
+        seq,
+        return_all_losses = False
+    ):
+        batch, device = seq.shape[0], seq.device
+
+        seq, labels = seq[:, :-1], seq[:, 1:]
+
+        encoder_mask = seq != self.pad_id
+
+        tokens = self.token_emb(seq)
+
+        # decoder head
+
+        tokens = self.decoder_head(tokens)
+
+        # get latent Z
+
+        latents, kl_loss = self.encode_to_latents(tokens, mask = encoder_mask, return_kl_loss = True)
+
+        condition = self.from_latent_to_condition(latents)
+
+        # decoder tail
+
+        tokens = self.decoder_tail(tokens)
+
+        # cross entropy loss
+
+        logits = self.token_unembed(tokens)
+
+        ar_loss = F.cross_entropy(
+            rearrange(logits, 'b n l -> b l n'),
+            labels,
+            ignore_index = self.pad_id
+        )
+
+        # return losses
+
+        total_loss = (
+            ar_loss +
+            kl_loss * self.kl_loss_weight
+        )
+
+        if not return_all_losses:
+            return total_loss
+
+        losses = (ar_loss, kl_loss)
+
+        return total_loss, losses