vector-quantize-pytorch 1.25.0__tar.gz → 1.25.2__tar.gz

{vector_quantize_pytorch-1.25.0 → vector_quantize_pytorch-1.25.2}/PKG-INFO

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

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

{vector_quantize_pytorch-1.25.0 → vector_quantize_pytorch-1.25.2}/vector_quantize_pytorch/binary_mapper.py

@@ -46,7 +46,8 @@ class BinaryMapper(Module):
     def __init__(
         self,
         bits = 1,
-        kl_loss_threshold = NAT # 1 bit
+        kl_loss_threshold = NAT, # 1 bit
+        deterministic_on_eval = False
     ):
         super().__init__()
 
@@ -64,14 +65,21 @@ class BinaryMapper(Module):
         self.kl_loss_threshold = kl_loss_threshold
         self.register_buffer('zero', tensor(0.), persistent = False)
 
+        # eval behavior
+
+        self.deterministic_on_eval = deterministic_on_eval
+
     def forward(
         self,
         logits,
         temperature = 1.,
         straight_through = None,
         calc_aux_loss = None,
-        return_indices = False
+        deterministic = None,
+        return_indices = False,
     ):
+        deterministic = default(deterministic, self.deterministic_on_eval and not self.training)
+
         straight_through = default(straight_through, self.training)
         calc_aux_loss = default(calc_aux_loss, self.training)
 
@@ -87,7 +95,11 @@ class BinaryMapper(Module):
 
         # sampling
 
-        sampled_bits = (torch.rand_like(logits) <= prob_for_sample).long()
+        if not deterministic:
+            sampled_bits = prob_for_sample.bernoulli().long()
+        else:
+            sampled_bits = (prob_for_sample > 0.5).long()
+
         indices = (self.power_two * sampled_bits).sum(dim = -1)
 
         one_hot = F.one_hot(indices, self.num_codes).float()
@@ -143,3 +155,8 @@ if __name__ == '__main__':
     assert sparse_one_hot.shape == (3, 4, 2 ** 8)
     assert indices.shape == (3, 4)
     assert aux_loss.numel() == 1
+
+    binary_mapper.eval()
+    sparse_one_hot1, _ = binary_mapper(logits, deterministic = True)
+    sparse_one_hot2, _ = binary_mapper(logits, deterministic = True)
+    assert torch.allclose(sparse_one_hot1, sparse_one_hot2)