final tweak to lfq

Author: lucidrains

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'vector_quantize_pytorch',
   packages = find_packages(),
-  version = '1.9.5',
+  version = '1.9.6',
   license='MIT',
   description = 'Vector Quantization - Pytorch',
   long_description_content_type = 'text/markdown',

vector_quantize_pytorch/lookup_free_quantization.py

@@ -53,7 +53,8 @@ def __init__(
         dim = None,
         codebook_size = None,
         entropy_loss_weight = 0.1,
-        diversity_gamma = 2.5
+        diversity_gamma = 2.5,
+        straight_through_activation = nn.Tanh()
     ):
         super().__init__()
 
@@ -73,6 +74,10 @@ def __init__(
         self.dim = dim
         self.codebook_dim = codebook_dim
 
+        # straight through activation
+
+        self.activation = straight_through_activation
+
         # entropy aux loss related weights
 
         self.diversity_gamma = diversity_gamma
@@ -122,7 +127,7 @@ def forward(
 
         is_img_or_video = x.ndim >= 4
 
-        # rearrange if image or video into (batch, seq, dimension)
+        # standardize image or video into (batch, seq, dimension)
 
         if is_img_or_video:
             x = rearrange(x, 'b d ... -> b ... d')
@@ -137,10 +142,10 @@ def forward(
         ones = torch.ones_like(x)
         quantized = torch.where(x > 0, ones, -ones)
 
-        # use straight-through gradients with tanh if training
+        # use straight-through gradients with tanh (or custom activation fn) if training
 
         if self.training:
-            x = torch.tanh(x * inv_temperature)
+            x = self.activation(x * inv_temperature)
             x = x - x.detach() + quantized
         else:
             x = quantized
@@ -181,6 +186,4 @@ def forward(
 
             indices = unpack_one(indices, ps, 'b *')
 
-        # bits to decimal for the codebook indices
-
         return Return(x, indices, entropy_aux_loss)