need to account for multi-headedness (multiple codebooks). cite einops for immeasurable time saved

Author: lucidrains

README.md

@@ -389,3 +389,13 @@ if __name__ == '__main__':
     organization = {PMLR}
 }
 ```
+
+```bibtex
+@inproceedings{rogozhnikov2022einops,
+    title   = {Einops: Clear and Reliable Tensor Manipulations with Einstein-like Notation},
+    author  = {Alex Rogozhnikov},
+    booktitle = {International Conference on Learning Representations},
+    year    = {2022},
+    url     = {https://openreview.net/forum?id=oapKSVM2bcj}
+}
+```

setup.py

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

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -323,11 +323,11 @@ def update_affine(self, data, embed):
 
         var_fn = partial(torch.var, unbiased = False)
 
-        self.update_with_decay('batch_mean', reduce(data, '... d -> d', 'mean'), self.affine_param_batch_decay)
-        self.update_with_decay('batch_variance', reduce(data, '... d -> d', var_fn), self.affine_param_batch_decay)
+        self.update_with_decay('batch_mean', reduce(data, 'h ... d -> h 1 d', 'mean'), self.affine_param_batch_decay)
+        self.update_with_decay('batch_variance', reduce(data, 'h ... d -> h 1 d', var_fn), self.affine_param_batch_decay)
 
-        self.update_with_decay('codebook_mean', reduce(embed, '... d -> d', 'mean'), self.affine_param_codebook_decay)
-        self.update_with_decay('codebook_variance', reduce(embed, '... d -> d', var_fn), self.affine_param_codebook_decay)
+        self.update_with_decay('codebook_mean', reduce(embed, 'h ... d -> h 1 d', 'mean'), self.affine_param_codebook_decay)
+        self.update_with_decay('codebook_variance', reduce(embed, 'h ... d -> h 1 d', var_fn), self.affine_param_codebook_decay)
 
     def replace(self, batch_samples, batch_mask):
         for ind, (samples, mask) in enumerate(zip(batch_samples.unbind(dim = 0), batch_mask.unbind(dim = 0))):