x-transformers 2.7.5__py3-none-any.whl → 2.8.0__py3-none-any.whl

x_transformers/gpt_vae.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+# applying the cvae + detr design from ACT (Zhou et al.) to GPT
+# for steering, diversity rlvr, map-elites in epo, and other possibilities
+
+import torch
+from torch import nn, Tensor, is_tensor, tensor
+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 AutoregressiveWrapper
+
+from einops.layers.torch import Rearrange
+from einops import rearrange, reduce, repeat
+
+# helper functions
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+# classes
+
+class GPTVAE(Module):
+    def __init__(
+        self,
+        *,
+        num_tokens,
+        dim,
+        depth,
+        enc_depth,
+        max_seq_len,
+        dim_latent = None,
+        attn_dim_head = 64,
+        heads = 8,
+        enc_kwargs: dict = dict(),
+        dec_kwargs: dict = dict(),
+        vae_kl_loss_weight = 1.,
+        latents_dropout_prob = 0.5, # what percentage of the time to dropout the latents completely
+        pad_id = -1,
+        **kwargs
+    ):
+        super().__init__()
+        dim_latent = default(dim_latent, dim)
+
+        self.encoder = TransformerWrapper(
+            num_tokens = num_tokens,
+            max_seq_len = max_seq_len + 1,
+            return_only_embed = True,
+            average_pool_embed = True,
+            attn_layers = Encoder(
+                dim = dim,
+                depth = enc_depth,
+                attn_dim_head = attn_dim_head,
+                heads = heads,
+                **kwargs,
+                **enc_kwargs
+            ),
+        )
+
+        self.to_latent_mean_log_variance = nn.Sequential(
+            nn.Linear(dim, dim_latent * 2),
+            Rearrange('b (two d) -> two b 1 d', two = 2)
+        )
+
+        self.from_latent_to_prepend_token = nn.Linear(dim_latent, dim)
+
+        self.decoder = TransformerWrapper(
+            num_tokens = num_tokens,
+            max_seq_len = max_seq_len,
+            attn_layers = Decoder(
+                dim = dim,
+                depth = depth,
+                attn_dim_head = attn_dim_head,
+                heads = heads,
+                **kwargs,
+                **dec_kwargs
+            ),
+        )
+
+        self.ar_wrapped_decoder = AutoregressiveWrapper(self.decoder, ignore_index = pad_id)
+
+        self.pad_id = pad_id
+
+        # loss weights - vae kl loss
+
+        self.vae_kl_loss_weight = vae_kl_loss_weight
+
+        self.latents_dropout = nn.Dropout(latents_dropout_prob)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def encode_to_latents(
+        self,
+        seq,
+        return_mean_log_var = False
+    ):
+        mask = seq != self.pad_id
+        pooled = self.encoder(seq, mask = mask)
+
+        latents_mean, latents_log_var = self.to_latent_mean_log_variance(pooled)
+        latents_std = (0.5 * latents_log_var).exp()
+
+        # reparam trick
+
+        latents = latents_mean + latents_std * torch.randn_like(latents_mean)
+
+        if not return_mean_log_var:
+            return latents
+
+        return latents, (latents_mean, latents_log_var)
+
+    @torch.no_grad()
+    def generate(
+        self,
+        prompts,
+        seq_len,
+        latents = None,
+        **generate_kwargs
+    ):
+        assert prompts.ndim in {1, 2}
+        batch = prompts.shape[0] if prompts.ndim == 2 else 1
+
+        # prepend embeds
+
+        prepend_embeds = 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)
+
+            prepend_embeds = self.from_latent_to_prepend_token(latents)
+
+        if exists(prepend_embeds):
+            prepend_embeds = rearrange(prepend_embeds, 'b d -> b 1 d')
+
+        # generated
+
+        generated = self.ar_wrapped_decoder.generate(
+            prompts,
+            seq_len,
+            prepend_embeds = prepend_embeds,
+            **generate_kwargs
+        )
+
+        return generated
+
+    def forward(
+        self,
+        seq,
+        return_all_losses = False
+    ):
+        batch, device = seq.shape[0], seq.device
+
+        latents, (latents_mean, latents_log_var) = self.encode_to_latents(seq, return_mean_log_var = True)
+
+        dropped_latents = ~self.latents_dropout(torch.ones((batch,), device = device)).bool()
+
+        prepend_embeds = self.from_latent_to_prepend_token(latents)
+
+        ar_loss = self.ar_wrapped_decoder(
+            seq,
+            prepend_embeds = prepend_embeds,
+            seq_start_pos = dropped_latents.long() # sequence starts at 1 and does not attend to the first style latent
+        )
+
+        # vae kl loss
+
+        vae_kl_loss = (
+            latents_log_var.exp()
+            + latents_mean.square()
+            - latents_log_var
+            - 1.
+        ).sum(dim = -1).mean()
+
+        # return losses
+
+        total_loss = (
+            ar_loss +
+            vae_kl_loss * self.vae_kl_loss_weight
+        )
+
+        if not return_all_losses:
+            return total_loss
+
+        losses = (ar_loss, vae_kl_loss)
+
+        return total_loss, losses

x_transformers/x_transformers.py

@@ -2469,12 +2469,15 @@ class AttentionLayers(Module):
     ):
         # pairs up the attention intermediates with each attention module and does qk clip proposed by kimi team
 
-        for (_, layer, _), layer_type, attn_inter in zip(self.layers, self.layer_types, intermediates.attn_intermediates):
+        layer_and_layer_types = (self.layers, self.layer_types)
 
-            if layer_type not in ('a', 'c'):
-                continue
+        attn_layers = [layer for (_, layer, _), layer_type in zip(self.layers, self.layer_types) if layer_type in ('a', 'c')]
+        attn_intermeds = intermediates.attn_intermediates
+
+        assert len(attn_layers) == len(attn_intermeds)
 
-            layer.qk_clip_(attn_inter, tau = tau)
+        for attn_layer, attn_inter in zip(attn_layers, attn_intermeds):
+            attn_layer.qk_clip_(attn_inter, tau = tau)
 
     def forward(
         self,

{x_transformers-2.7.5.dist-info → x_transformers-2.8.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: x-transformers
-Version: 2.7.5
+Version: 2.8.0
 Summary: X-Transformers
 Project-URL: Homepage, https://pypi.org/project/x-transformers/
 Project-URL: Repository, https://github.com/lucidrains/x-transformers
@@ -2540,4 +2540,16 @@ ids_out, num_out, is_number_mask = model.generate(start_ids, start_nums, 17)
 }
 ```
 
+```bibtex
+@misc{zhao2023learningfinegrainedbimanualmanipulation,
+    title   = {Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware}, 
+    author  = {Tony Z. Zhao and Vikash Kumar and Sergey Levine and Chelsea Finn},
+    year    = {2023},
+    eprint  = {2304.13705},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.RO},
+    url     = {https://arxiv.org/abs/2304.13705}, 
+}
+```
+
 *solve intelligence... then use that to solve everything else.* - Demis Hassabis

{x_transformers-2.7.5.dist-info → x_transformers-2.8.0.dist-info}/RECORD

@@ -5,14 +5,15 @@ x_transformers/belief_state_wrapper.py,sha256=YLUMk6t2MhFBEw5lHDDHJHcoCxTIkHvxTN
 x_transformers/continuous.py,sha256=WwpQCjyVY4PtuEAOFY68zqgklbF9I7AL5w6874YlDe8,13249
 x_transformers/dpo.py,sha256=xt4OuOWhU8pN3OKN2LZAaC2NC8iiEnchqqcrPWVqf0o,3521
 x_transformers/entropy_based_tokenizer.py,sha256=F2lO8-v3aLIcVDVNhu7RR-UtRdlmaaYJzBK9m7OnLE8,5018
+x_transformers/gpt_vae.py,sha256=yqL1K2yJ6RSP_MC6XSHI3hjiUnaptddg6CUnbEX4Bsk,5281
 x_transformers/multi_input.py,sha256=tCh-fTJDj2ib4SMGtsa-AM8MxKzJAQSwqAXOu3HU2mg,9252
 x_transformers/neo_mlp.py,sha256=XCNnnop9WLarcxap1kGuYc1x8GHvwkZiDRnXOxSl3Po,3452
 x_transformers/nonautoregressive_wrapper.py,sha256=hMQqNimGtchNIe13cR5LZule1V7I1qM5LmY8VQfVdnA,11698
 x_transformers/up_wrapper.py,sha256=YC2LN14_7Xx9Wtiek2rtEJ_qHqdfSmKlh3d7Cgxwd80,7073
-x_transformers/x_transformers.py,sha256=xaGBkYCy6CqL0q9icWmL_WzCeU6ZztEYEkMtN71L2z4,124576
+x_transformers/x_transformers.py,sha256=odnCZAKZKrQLXmpaWhiPVB5elGjt8kerDbO3-yeC-60,124764
 x_transformers/xl_autoregressive_wrapper.py,sha256=CvZMJ6A6PA-Y_bQAhnORwjJBSl6Vjq2IdW5KTdk8NI8,4195
 x_transformers/xval.py,sha256=AwwYUm8yDAtKQyKJDIhYMsiLTJ_skh3scUFMjp5sda8,8597
-x_transformers-2.7.5.dist-info/METADATA,sha256=m6f4PIgJFKKWlsGAydi_Bg5-7-0IRlor0pRY_zBh5s8,93739
-x_transformers-2.7.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-x_transformers-2.7.5.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
-x_transformers-2.7.5.dist-info/RECORD,,
+x_transformers-2.8.0.dist-info/METADATA,sha256=jPo0ZPhD1d_aocaDqFYWXA7EXPAcxWeUYNDzKpY1yi8,94136
+x_transformers-2.8.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+x_transformers-2.8.0.dist-info/licenses/LICENSE,sha256=As9u198X-U-vph5noInuUfqsAG2zX_oXPHDmdjwlPPY,1066
+x_transformers-2.8.0.dist-info/RECORD,,