add ability to easily get codes for multi-headed vq

Author: lucidrains

setup.py

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

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -4,7 +4,7 @@
 import torch.distributed as distributed
 from torch.cuda.amp import autocast
 
-from einops import rearrange, repeat
+from einops import rearrange, repeat, pack, unpack
 from contextlib import contextmanager
 
 def exists(val):
@@ -470,6 +470,7 @@ def __init__(
         sync_codebook = False
     ):
         super().__init__()
+        self.dim = dim
         self.heads = heads
         self.separate_codebook_per_head = separate_codebook_per_head
 
@@ -518,6 +519,25 @@ def codebook(self):
 
         return rearrange(codebook, '1 ... -> ...')
 
+    def get_codes_from_indices(self, indices):
+        codebook = self.codebook
+        is_multiheaded = codebook.ndim > 2
+
+        if not is_multiheaded:
+            codes = codebook[indices]
+            return rearrange(codes, '... h d -> ... (h d)')
+
+        indices, ps = pack([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])
+
+        codes = codebook.gather(2, indices)
+        codes = rearrange(codes, 'b h n d -> b n (h d)')
+        codes, = unpack(codes, ps, 'b * d')
+        return codes
+
     def forward(
         self,
         x,