address #233

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.27.6"
+version = "1.27.7"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

tests/test_beam.py

@@ -1,3 +1,6 @@
+import pytest
+param = pytest.mark.parametrize
+
 import torch
 from vector_quantize_pytorch import VectorQuantize
 
@@ -43,20 +46,24 @@ def test_topk_and_manual_ema_update():
     assert torch.allclose(vq1._codebook.embed_avg, vq2._codebook.embed_avg)
     assert torch.allclose(vq1.codebook, vq2.codebook)
 
-def test_beam_search():
+@param('codebook_dim', (256, 128))
+def test_beam_search(
+    codebook_dim
+):
     import torch
     from vector_quantize_pytorch import ResidualVQ
 
     residual_vq = ResidualVQ(
         dim = 256,
+        codebook_dim = codebook_dim,
         num_quantizers = 8,      # specify number of quantizers
         codebook_size = 1024,    # codebook size
         quantize_dropout = True,
         beam_size = 2,
         eval_beam_size = 3
     )
 
-    x = torch.randn(1, 1024, 256)
+    x = torch.randn(1, 1024, 256).requires_grad_()
 
     for _ in range(5):
         quantized, indices, commit_loss = residual_vq(x)

vector_quantize_pytorch/residual_vq.py

@@ -191,6 +191,7 @@ def __init__(
 
         codebook_dim = default(codebook_dim, dim)
         codebook_input_dim = codebook_dim * heads
+        self.codebook_dim = codebook_dim
 
         requires_projection = codebook_input_dim != dim
         self.project_in = nn.Linear(dim, codebook_input_dim) if requires_projection else nn.Identity()
@@ -223,6 +224,7 @@ def __init__(
         self.num_quantizers = num_quantizers
 
         self.codebook_sizes = codebook_sizes
+
         self.uniform_codebook_size = len(unique(codebook_sizes)) == 1
 
         # define vq across layers
@@ -287,10 +289,6 @@ def __init__(
     @property
     def codebook_size(self):
         return self.layers[0].codebook_size
-    
-    @property
-    def codebook_dim(self):
-        return self.layers[0].codebook_dim
 
     @property
     def codebooks(self):
@@ -423,7 +421,7 @@ def forward(
 
         # save all inputs across layers, for use during expiration at end under shared codebook setting, or ema update during beam search
 
-        all_residuals = torch.empty((*input_shape[:-1], 0, input_shape[-1]), dtype = residual.dtype, device = device)
+        all_residuals = torch.empty((*input_shape[:-1], 0, self.codebook_dim), dtype = residual.dtype, device = device)
 
         # maybe prepare beam search
 

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -1134,6 +1134,11 @@ def forward(
 
             commit_quantize = maybe_detach(quantize)
 
+            # maybe expand input if returning topk codes
+
+            if exists(topk):
+                x = repeat(x, '... d -> ... k d', k = topk)
+
             # spare rotation trick calculation if inputs do not need gradients
 
             if input_requires_grad: