vector-quantize-pytorch 1.22.4__tar.gz → 1.22.5__tar.gz

{vector_quantize_pytorch-1.22.4 → vector_quantize_pytorch-1.22.5}/PKG-INFO

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

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

{vector_quantize_pytorch-1.22.4 → vector_quantize_pytorch-1.22.5}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -138,11 +138,15 @@ def sample_multinomial(total_count, probs):
     remainder = probs.new_ones(())
     sample = torch.empty_like(probs, dtype = torch.long)
 
-    for i, p in enumerate(probs):
-        s = torch.binomial(total_count, p / remainder)
+    num_probs = len(probs)
+
+    for i, prob in enumerate(probs):
+        is_last = i == (num_probs - 1)
+
+        s = torch.binomial(total_count, prob / remainder) if not is_last else total_count
         sample[i] = s
         total_count -= s
-        remainder -= p
+        remainder -= prob
 
     assert total_count == 0, f'invalid total count {total_count}'