start leveraging einx get_at for extra clarity

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '1.12.12',
+  version = '1.12.14',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   long_description_content_type = 'text/markdown',
@@ -18,6 +18,7 @@
   ],
   install_requires=[
     'einops>=0.7.0',
+    'einx',
     'torch'
   ],
   classifiers=[

vector_quantize_pytorch/residual_lfq.py

@@ -12,6 +12,8 @@
 
 from einops import rearrange, repeat, reduce, pack, unpack
 
+from einx import get_at
+
 # helper functions
 
 def exists(val):
@@ -92,17 +94,12 @@ def get_codes_from_indices(self, indices):
             assert self.quantize_dropout > 0., 'quantize dropout must be greater than 0 if you wish to reconstruct from a signal with less fine quantizations'
             indices = F.pad(indices, (0, self.num_quantizers - quantize_dim), value = -1)
 
-        # get ready for gathering
-
-        codebooks = repeat(self.codebooks, 'q c d -> q b c d', b = batch)
-        gather_indices = repeat(indices, 'b n q -> q b n d', d = codebooks.shape[-1])
-
         # take care of quantizer dropout
 
-        mask = gather_indices == -1.
-        gather_indices = gather_indices.masked_fill(mask, 0) # have it fetch a dummy code to be masked out later
+        mask = indices == -1.
+        indices = indices.masked_fill(mask, 0) # have it fetch a dummy code to be masked out later
 
-        all_codes = codebooks.gather(2, gather_indices) # gather all codes
+        all_codes = get_at('q [c] d, b n q -> q b n d', self.codebooks, indices)
 
         # mask out any codes that were dropout-ed
 

vector_quantize_pytorch/residual_vq.py

@@ -10,6 +10,8 @@
 
 from einops import rearrange, repeat, reduce, pack, unpack
 
+from einx import get_at
+
 # helper functions
 
 def exists(val):
@@ -94,17 +96,12 @@ def get_codes_from_indices(self, indices):
             assert self.quantize_dropout > 0., 'quantize dropout must be greater than 0 if you wish to reconstruct from a signal with less fine quantizations'
             indices = F.pad(indices, (0, self.num_quantizers - quantize_dim), value = -1)
 
-        # get ready for gathering
-
-        codebooks = repeat(self.codebooks, 'q c d -> q b c d', b = batch)
-        gather_indices = repeat(indices, 'b n q -> q b n d', d = codebooks.shape[-1])
-
         # take care of quantizer dropout
 
-        mask = gather_indices == -1.
-        gather_indices = gather_indices.masked_fill(mask, 0) # have it fetch a dummy code to be masked out later
+        mask = indices == -1.
+        indices = indices.masked_fill(mask, 0) # have it fetch a dummy code to be masked out later
 
-        all_codes = codebooks.gather(2, gather_indices) # gather all codes
+        all_codes = get_at('q [c] d, b n q -> q b n d', self.codebooks, indices)
 
         # mask out any codes that were dropout-ed