offer get_output_from_indices for VectorQuantize, ResidualVQ, and ResidualLFQ - which takes into account linear projection from codebook dimension back to the input dimension

Author: lucidrains

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '1.10.1',
+  version = '1.10.2',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   long_description_content_type = 'text/markdown',

vector_quantize_pytorch/residual_lfq.py

@@ -10,7 +10,7 @@
 
 from vector_quantize_pytorch.lookup_free_quantization import LFQ
 
-from einops import rearrange, repeat, pack, unpack
+from einops import rearrange, repeat, reduce, pack, unpack
 
 # helper functions
 
@@ -111,6 +111,11 @@ def get_codes_from_indices(self, indices):
 
         return all_codes
 
+    def get_output_from_indices(self, indices):
+        codes = self.get_codes_from_indices(indices)
+        codes_summed = reduce(codes, 'q ... -> ...', 'sum')
+        return self.project_out(codes_summed)
+
     def forward(
         self,
         x,

vector_quantize_pytorch/residual_vq.py

@@ -8,7 +8,7 @@
 import torch.nn.functional as F
 from vector_quantize_pytorch.vector_quantize_pytorch import VectorQuantize
 
-from einops import rearrange, repeat, pack, unpack
+from einops import rearrange, repeat, reduce, pack, unpack
 
 # helper functions
 
@@ -113,6 +113,11 @@ def get_codes_from_indices(self, indices):
 
         return all_codes
 
+    def get_output_from_indices(self, indices):
+        codes = self.get_codes_from_indices(indices)
+        codes_summed = reduce(codes, 'q ... -> ...', 'sum')
+        return self.project_out(codes_summed)
+
     def forward(
         self,
         x,

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -834,6 +834,10 @@ def get_codes_from_indices(self, indices):
         codes = unpack_one(codes, ps, 'b * d')
         return codes
 
+    def get_output_from_indices(self, indices):
+        codes = self.get_codes_from_indices(indices)
+        return self.project_out(codes)
+
     def forward(
         self,
         x,
@@ -964,7 +968,7 @@ def calculate_ce_loss(codes):
             embed_ind = rearrange(embed_ind, 'b (h w) ... -> b h w ...', h = height, w = width)
 
         if only_one:
-            embed_ind = rearrange(embed_ind, 'b 1 -> b')
+            embed_ind = rearrange(embed_ind, 'b 1 ... -> b ...')
 
         # aggregate loss