vector-quantize-pytorch 1.15.2__tar.gz → 1.15.4__tar.gz

{vector_quantize_pytorch-1.15.2 → vector_quantize_pytorch-1.15.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: vector-quantize-pytorch
-Version: 1.15.2
+Version: 1.15.4
 Summary: Vector Quantization - Pytorch
 Project-URL: Homepage, https://pypi.org/project/vector-quantize-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/vector-quantizer-pytorch
@@ -35,7 +35,7 @@ Classifier: Programming Language :: Python :: 3.6
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Requires-Dist: einops>=0.8.0
-Requires-Dist: einx>=0.2.2
+Requires-Dist: einx>=0.3.0
 Requires-Dist: torch>=2.0
 Provides-Extra: examples
 Requires-Dist: torchvision; extra == 'examples'
@@ -723,3 +723,12 @@ assert loss.item() >= 0
     primaryClass = {cs.LG}
 }
 ```
+
+```bibtex
+@inproceedings{Irie2023SelfOrganisingND,
+    title   = {Self-Organising Neural Discrete Representation Learning \`a la Kohonen},
+    author  = {Kazuki Irie and R'obert Csord'as and J{\"u}rgen Schmidhuber},
+    year    = {2023},
+    url     = {https://api.semanticscholar.org/CorpusID:256901024}
+}
+```

{vector_quantize_pytorch-1.15.2 → vector_quantize_pytorch-1.15.4}/README.md

@@ -679,3 +679,12 @@ assert loss.item() >= 0
     primaryClass = {cs.LG}
 }
 ```
+
+```bibtex
+@inproceedings{Irie2023SelfOrganisingND,
+    title   = {Self-Organising Neural Discrete Representation Learning \`a la Kohonen},
+    author  = {Kazuki Irie and R'obert Csord'as and J{\"u}rgen Schmidhuber},
+    year    = {2023},
+    url     = {https://api.semanticscholar.org/CorpusID:256901024}
+}
+```

{vector_quantize_pytorch-1.15.2 → vector_quantize_pytorch-1.15.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.15.2"
+version = "1.15.4"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -25,7 +25,7 @@ classifiers=[
 dependencies = [
     "torch>=2.0",
     "einops>=0.8.0",
-    "einx>=0.2.2",
+    "einx>=0.3.0",
 ]
 
 [project.urls]

{vector_quantize_pytorch-1.15.2 → vector_quantize_pytorch-1.15.4}/tests/test_readme.py

@@ -37,6 +37,31 @@ def test_vq_eval():
     quantized, indices, commit_loss = vq(x)
     assert torch.allclose(quantized, vq.get_output_from_indices(indices))
 
+def test_vq_mask():
+    from vector_quantize_pytorch import VectorQuantize
+
+    vq = VectorQuantize(
+        dim = 256,
+        codebook_size = 512,     # codebook size
+        decay = 1.,             # the exponential moving average decay, lower means the dictionary will change faster
+        commitment_weight = 1.   # the weight on the commitment loss
+    )
+
+    x = torch.randn(1, 1024, 256)
+    lens = torch.full((1,), 512)
+
+    vq.train()
+
+    quantized, indices, commit_loss = vq(x[:, :512])
+    mask_quantized, mask_indices, mask_commit_loss = vq(x, lens = lens)
+
+    assert torch.allclose(commit_loss, mask_commit_loss)
+    assert torch.allclose(quantized, mask_quantized[:, :512])
+    assert torch.allclose(indices, mask_indices[:, :512])
+
+    assert torch.allclose(mask_quantized[:, 512:], x[:, 512:])
+    assert (mask_indices[:, 512:] == -1).all()
+
 def test_residual_vq():
     from vector_quantize_pytorch import ResidualVQ
 

{vector_quantize_pytorch-1.15.2 → vector_quantize_pytorch-1.15.4}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -3,12 +3,13 @@ from collections import namedtuple
 
 import torch
 from torch.nn import Module
-from torch import nn, einsum
+from torch import nn, einsum, Tensor
 import torch.nn.functional as F
 import torch.distributed as distributed
 from torch.optim import Optimizer
 from torch.cuda.amp import autocast
 
+import einx
 from einops import rearrange, repeat, reduce, pack, unpack
 
 from typing import Callable
@@ -63,6 +64,10 @@ def pack_one(t, pattern):
 def unpack_one(t, ps, pattern):
     return unpack(t, ps, pattern)[0]
 
+def lens_to_mask(lens, max_length):
+    seq = torch.arange(max_length, device = lens.device)
+    return seq < lens[:, None]
+
 def uniform_init(*shape):
     t = torch.empty(shape)
     nn.init.kaiming_uniform_(t)
@@ -312,7 +317,7 @@ class EuclideanCodebook(Module):
         self.all_reduce_fn = distributed.all_reduce if use_ddp else noop
 
         self.register_buffer('initted', torch.Tensor([not kmeans_init]))
-        self.register_buffer('cluster_size', torch.zeros(num_codebooks, codebook_size))
+        self.register_buffer('cluster_size', torch.ones(num_codebooks, codebook_size))
         self.register_buffer('embed_avg', embed.clone())
 
         self.learnable_codebook = learnable_codebook
@@ -582,7 +587,7 @@ class CosineSimCodebook(Module):
         self.all_reduce_fn = distributed.all_reduce if use_ddp else noop
 
         self.register_buffer('initted', torch.Tensor([not kmeans_init]))
-        self.register_buffer('cluster_size', torch.zeros(num_codebooks, codebook_size))
+        self.register_buffer('cluster_size', torch.ones(num_codebooks, codebook_size))
         self.register_buffer('embed_avg', embed.clone())
 
         self.learnable_codebook = learnable_codebook
@@ -897,12 +902,22 @@ class VectorQuantize(Module):
         x,
         indices = None,
         mask = None,
+        lens = None,
         sample_codebook_temp = None,
         freeze_codebook = False,
         return_loss_breakdown = False,
     ):
         orig_input = x
 
+        # handle masking, either passed in as `mask` or `lens`
+
+        assert not (exists(mask) and exists(lens))
+
+        if exists(lens):
+            mask = lens_to_mask(lens, x.shape[1])
+
+        # handle one token given
+
         only_one = x.ndim == 2
 
         if only_one:
@@ -917,6 +932,7 @@ class VectorQuantize(Module):
         # rearrange inputs
 
         if self.accept_image_fmap:
+            assert not exists(mask)
             height, width = x.shape[-2:]
             x = rearrange(x, 'b c h w -> b (h w) c')
 
@@ -1117,12 +1133,20 @@ class VectorQuantize(Module):
         # if masking, only return quantized for where mask has True
 
         if exists(mask):
-            quantize = torch.where(
-                rearrange(mask, '... -> ... 1'),
+            quantize = einx.where(
+                'b n, b n d, b n d -> b n d',
+                mask,
                 quantize,
                 orig_input
             )
 
+            embed_ind = einx.where(
+                'b n, b n ..., -> b n ...',
+                mask,
+                embed_ind,
+                -1
+            )
+
         if not return_loss_breakdown:
             return quantize, embed_ind, loss