vector-quantize-pytorch 1.16.1__tar.gz → 1.16.2__tar.gz

{vector_quantize_pytorch-1.16.1 → vector_quantize_pytorch-1.16.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: vector-quantize-pytorch
-Version: 1.16.1
+Version: 1.16.2
 Summary: Vector Quantization - Pytorch
 Project-URL: Homepage, https://pypi.org/project/vector-quantize-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/vector-quantizer-pytorch

{vector_quantize_pytorch-1.16.1 → vector_quantize_pytorch-1.16.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.16.1"
+version = "1.16.2"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{vector_quantize_pytorch-1.16.1 → vector_quantize_pytorch-1.16.2}/tests/test_readme.py

@@ -64,9 +64,11 @@ def test_vq_mask():
 
 @pytest.mark.parametrize('implicit_neural_codebook', (True, False))
 @pytest.mark.parametrize('use_cosine_sim', (True, False))
+@pytest.mark.parametrize('train', (True, False))
 def test_residual_vq(
     implicit_neural_codebook,
-    use_cosine_sim
+    use_cosine_sim,
+    train
 ):
     from vector_quantize_pytorch import ResidualVQ
 
@@ -80,14 +82,9 @@ def test_residual_vq(
 
     x = torch.randn(1, 256, 32)
 
-    quantized, indices, commit_loss = residual_vq(x)
-    quantized, indices, commit_loss, all_codes = residual_vq(x, return_all_codes = True)
-
-    # test eval mode and `get_output_from_indices`
-
-    residual_vq.eval()
-    quantized, indices, commit_loss = residual_vq(x)
+    residual_vq.train(train)
 
+    quantized, indices, commit_loss = residual_vq(x, freeze_codebook = train and not implicit_neural_codebook)
     quantized_out = residual_vq.get_output_from_indices(indices)
     assert torch.allclose(quantized, quantized_out, atol = 1e-6)
 

{vector_quantize_pytorch-1.16.1 → vector_quantize_pytorch-1.16.2}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -751,7 +751,7 @@ class CosineSimCodebook(Module):
             embed_normalized = self.embed_avg / rearrange(cluster_size, '... -> ... 1')
             embed_normalized = l2norm(embed_normalized)
 
-            self.embed.data.copy_(l2norm(embed_normalized))
+            self.embed.data.copy_(embed_normalized)
             self.expire_codes_(x)
 
         if needs_codebook_dim: