vector-quantize-pytorch 1.22.7__tar.gz → 1.22.9__tar.gz

{vector_quantize_pytorch-1.22.7 → vector_quantize_pytorch-1.22.9}/PKG-INFO

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

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

{vector_quantize_pytorch-1.22.7 → vector_quantize_pytorch-1.22.9}/vector_quantize_pytorch/finite_scalar_quantization.py

@@ -137,22 +137,19 @@ class FSQ(Module):
 
     def quantize(self, z):
         """ Quantizes z, returns quantized zhat, same shape as z. """
+
         shape, device, noise_dropout, preserve_symmetry, half_width = z.shape[0], z.device, self.noise_dropout, self.preserve_symmetry, (self._levels // 2)
+        bound_fn = self.symmetry_preserving_bound if preserve_symmetry else self.bound
 
         # determine where to add a random offset elementwise
         # if using noise dropout
 
         if self.training and noise_dropout > 0.:
             offset_mask = torch.bernoulli(torch.full_like(z, noise_dropout)).bool()
-            offset = (torch.rand_like(z) - 0.5) / half_width
+            offset = torch.rand_like(z) - 0.5
             z = torch.where(offset_mask, z + offset, z)
 
-        if preserve_symmetry:
-            quantized = round_ste(self.symmetry_preserving_bound(z)) / half_width
-        else:
-            quantized = round_ste(self.bound(z)) / half_width
-
-        return quantized
+        return round_ste(bound_fn(z)) / half_width
 
     def _scale_and_shift(self, zhat_normalized):
         half_width = self._levels // 2