vector-quantize-pytorch 1.23.2__tar.gz → 1.23.4__tar.gz

{vector_quantize_pytorch-1.23.2 → vector_quantize_pytorch-1.23.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vector-quantize-pytorch
-Version: 1.23.2
+Version: 1.23.4
 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.23.2 → vector_quantize_pytorch-1.23.4}/pyproject.toml

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

{vector_quantize_pytorch-1.23.2 → vector_quantize_pytorch-1.23.4}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -273,12 +273,14 @@ def efficient_rotation_trick_transform(u, q, e):
     e = rearrange(e, 'b d -> b 1 d')
     w = l2norm(u + q, dim = 1).detach()
 
-    return (
+    out = (
         e -
         2 * (e @ rearrange(w, 'b d -> b d 1') @ rearrange(w, 'b d -> b 1 d')) +
         2 * (e @ rearrange(u, 'b d -> b d 1').detach() @ rearrange(q, 'b d -> b 1 d').detach())
     )
 
+    return rearrange(out, '... 1 -> ...')
+
 def rotate_to(src, tgt):
     # rotation trick STE (https://arxiv.org/abs/2410.06424) to get gradients through VQ layer.
     src, inverse = pack_one(src, '* d')
@@ -291,7 +293,7 @@ def rotate_to(src, tgt):
         safe_div(src, norm_src),
         safe_div(tgt, norm_tgt),
         src
-    ).squeeze()
+    )
 
     rotated = rotated_tgt * safe_div(norm_tgt, norm_src).detach()
 
@@ -896,7 +898,7 @@ class VectorQuantize(Module):
         stochastic_sample_codes = False,
         sample_codebook_temp = 1.,
         straight_through = False,
-        rotation_trick = True,  # Propagate grads through VQ layer w/ rotation trick: https://arxiv.org/abs/2410.06424 by @cfifty
+        rotation_trick = None,  # Propagate grads through VQ layer w/ rotation trick: https://arxiv.org/abs/2410.06424 by @cfifty
         sync_codebook = None,
         sync_affine_param = False,
         ema_update = True,
@@ -911,6 +913,8 @@ class VectorQuantize(Module):
         return_zeros_for_masked_padding = True
     ):
         super().__init__()
+        rotation_trick = default(rotation_trick, dim > 1) # only use rotation trick if feature dimension greater than 1
+
         self.dim = dim
         self.heads = heads
         self.separate_codebook_per_head = separate_codebook_per_head
@@ -1263,7 +1267,7 @@ class VectorQuantize(Module):
             # calculate codebook diversity loss (negative of entropy) if needed
 
             if self.has_codebook_diversity_loss:
-                prob = (-distances * self.codebook_diversity_temperature).softmax(dim = -1)
+                prob = (distances * self.codebook_diversity_temperature).softmax(dim = -1)
                 avg_prob = reduce(prob, '... n l -> n l', 'mean')
                 codebook_diversity_loss = -entropy(avg_prob).mean()