fix channel_last = False for VectorQuantize when using get_output_from_indices

Author: lucidrains

setup.py

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

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -826,17 +826,21 @@ def get_codes_from_indices(self, indices):
         is_multiheaded = codebook.ndim > 2
 
         if not is_multiheaded:
-            return codebook[indices]
+            codes = codebook[indices]
+        else:
+            indices, ps = pack_one(indices, 'b * h')
+            indices = rearrange(indices, 'b n h -> b h n')
+
+            indices = repeat(indices, 'b h n -> b h n d', d = codebook.shape[-1])
+            codebook = repeat(codebook, 'h n d -> b h n d', b = indices.shape[0])
 
-        indices, ps = pack_one(indices, 'b * h')
-        indices = rearrange(indices, 'b n h -> b h n')
+            codes = codebook.gather(2, indices)
+            codes = rearrange(codes, 'b h n d -> b n (h d)')
+            codes = unpack_one(codes, ps, 'b * d')
 
-        indices = repeat(indices, 'b h n -> b h n d', d = codebook.shape[-1])
-        codebook = repeat(codebook, 'h n d -> b h n d', b = indices.shape[0])
+        if not self.channel_last:
+            codes = rearrange(codes, 'b ... d -> b d ...')
 
-        codes = codebook.gather(2, indices)
-        codes = rearrange(codes, 'b h n d -> b n (h d)')
-        codes = unpack_one(codes, ps, 'b * d')
         return codes
 
     def get_output_from_indices(self, indices):