vector-quantize-pytorch 1.31.2__tar.gz → 1.31.4__tar.gz

{vector_quantize_pytorch-1.31.2 → vector_quantize_pytorch-1.31.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: vector-quantize-pytorch
-Version: 1.31.2
+Version: 1.31.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.31.2 → vector_quantize_pytorch-1.31.4}/pyproject.toml

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

{vector_quantize_pytorch-1.31.2 → vector_quantize_pytorch-1.31.4}/vector_quantize_pytorch/residual_vq.py

@@ -185,6 +185,7 @@ class ResidualVQ(Module):
         eval_beam_size = None,
         beam_score_quantizer_weights: list[float] | None = None,
         quant_grad_frac = 0.,
+        return_zeros_for_masked_padding = True,
         **vq_kwargs
     ):
         super().__init__()
@@ -202,6 +203,9 @@ class ResidualVQ(Module):
 
         self.accept_image_fmap = accept_image_fmap
 
+        self.return_zeros_for_masked_padding = return_zeros_for_masked_padding
+        vq_kwargs.update(return_zeros_for_masked_padding = return_zeros_for_masked_padding)
+
         self.implicit_neural_codebook = implicit_neural_codebook
 
         if implicit_neural_codebook:
@@ -397,6 +401,8 @@ class ResidualVQ(Module):
 
         input_shape, num_quant, quant_dropout_multiple_of, return_loss, device = x.shape, self.num_quantizers, self.quantize_dropout_multiple_of, exists(indices), x.device
 
+        orig_input = x
+
         beam_size = default(beam_size, self.beam_size if self.training else self.eval_beam_size)
 
         is_beam_search = exists(beam_size) and beam_size > 1
@@ -577,7 +583,7 @@ class ResidualVQ(Module):
 
             if exists(mask):
                 all_losses = einx.where('..., ... l,', mask, all_losses, 0.)
-                all_losses = reduce(all_losses, '... l -> l', 'sum') / mask.sum(dim = -1).clamp_min(1e-4)
+                all_losses = reduce(all_losses, '... l -> l', 'sum') / mask.sum().clamp_min(1e-4)
             else:
                 all_losses = reduce(all_losses, '... l -> l', 'mean')
 
@@ -609,6 +615,21 @@ class ResidualVQ(Module):
 
         quantized_out = self.project_out(quantized_out)
 
+        # if masking, only return quantized for where mask has True
+
+        if exists(mask):
+            masked_out_value = orig_input
+
+            if self.return_zeros_for_masked_padding:
+                masked_out_value = torch.zeros_like(quantized_out)
+
+            quantized_out = einx.where(
+                'b n, b n d, b n d -> b n d',
+                mask,
+                quantized_out,
+                masked_out_value
+            )
+
         # whether to early return the cross entropy loss
 
         if return_loss:

{vector_quantize_pytorch-1.31.2 → vector_quantize_pytorch-1.31.4}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -1154,6 +1154,7 @@ class VectorQuantize(Module):
 
         if need_transpose:
             x = rearrange(x, 'b d n -> b n d')
+            orig_input = rearrange(orig_input, 'b d n -> b n d')
 
         # project input
 
@@ -1295,7 +1296,7 @@ class VectorQuantize(Module):
             commit_loss = reduce(commit_loss, '... k d -> ... k', 'mean')
 
             if exists(mask):
-                commit_loss = einx.where('..., ... k, -> ... k', mask, commit_loss, 0.)
+                commit_loss = einx.where('b n, b n ... k, -> b n ... k', mask, commit_loss, 0.)
 
             loss = commit_loss * self.commitment_weight if self.has_commitment_loss else commit_loss
 
@@ -1327,7 +1328,7 @@ class VectorQuantize(Module):
                         commit_loss = reduce(commit_loss, '... k d -> ... k', 'mean')
 
                         if exists(mask):
-                            commit_loss = einx.where('..., ... k, -> ... k', mask, commit_loss, 0.)
+                            commit_loss = einx.where('b n, b n ... k, -> b n ... k', mask, commit_loss, 0.)
 
                     elif exists(mask):
                         # with variable lengthed sequences
@@ -1386,20 +1387,6 @@ class VectorQuantize(Module):
 
         quantize = self.project_out(quantize)
 
-        # rearrange quantized embeddings
-
-        if need_transpose:
-            quantize = rearrange(quantize, 'b n ... d -> b d n ...')
-
-        if self.accept_image_fmap:
-            quantize = rearrange(quantize, 'b (h w) ... c -> b c h w ...', h = height, w = width)
-
-        if self.accept_3d_fmap:
-            quantize = rearrange(quantize, "b (d h w) ... c -> b c d h w ...", d=depth, h=height, w=width)
-
-        if only_one:
-            quantize = rearrange(quantize, 'b 1 ... d -> b ... d')
-
         # if masking, only return quantized for where mask has True
 
         if exists(mask):
@@ -1408,8 +1395,11 @@ class VectorQuantize(Module):
             if self.return_zeros_for_masked_padding:
                 masked_out_value = torch.zeros_like(orig_input)
 
+            if is_topk:
+                masked_out_value = repeat(masked_out_value, '... d -> ... k d', k = topk)
+
             quantize = einx.where(
-                'b n, b n ... d, b n d -> b n ... d',
+                'b n, b n ... d, b n ... d -> b n ... d',
                 mask,
                 quantize,
                 masked_out_value
@@ -1422,6 +1412,20 @@ class VectorQuantize(Module):
                 -1
             )
 
+        # rearrange quantized embeddings
+
+        if need_transpose:
+            quantize = rearrange(quantize, 'b n ... d -> b d n ...')
+
+        if self.accept_image_fmap:
+            quantize = rearrange(quantize, 'b (h w) ... c -> b c h w ...', h = height, w = width)
+
+        if self.accept_3d_fmap:
+            quantize = rearrange(quantize, "b (d h w) ... c -> b c d h w ...", d=depth, h=height, w=width)
+
+        if only_one:
+            quantize = rearrange(quantize, 'b 1 ... d -> b ... d')
+
         if not return_loss_breakdown:
             return quantize, embed_ind, loss