handle only 1 embedding given, release 1.0.0, should be matured

Author: lucidrains

setup.py

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

vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -526,6 +526,11 @@ def forward(
         x,
         mask = None
     ):
+        only_one = x.ndim == 2
+
+        if only_one:
+            x = rearrange(x, 'b d -> b 1 d')
+
         shape, device, heads, is_multiheaded, codebook_size = x.shape, x.device, self.heads, self.heads > 1, self.codebook_size
 
         need_transpose = not self.channel_last and not self.accept_image_fmap
@@ -600,4 +605,8 @@ def forward(
             quantize = rearrange(quantize, 'b (h w) c -> b c h w', h = height, w = width)
             embed_ind = rearrange(embed_ind, 'b (h w) ... -> b h w ...', h = height, w = width)
 
+        if only_one:
+            quantize = rearrange(quantize, 'b 1 d -> b d')
+            embed_ind = rearrange(embed_ind, 'b 1 -> b')
+
         return quantize, embed_ind, loss