add ability to softclamp the initial input going into residual fsq

Author: lucidrains

pyproject.toml

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

vector_quantize_pytorch/residual_fsq.py

@@ -59,6 +59,7 @@ def __init__(
         quantize_dropout = False,
         quantize_dropout_cutoff_index = 0,
         quantize_dropout_multiple_of = 1,
+        soft_clamp_input_value = None,
         **kwargs
     ):
         super().__init__()
@@ -73,6 +74,12 @@ def __init__(
         self.is_channel_first = is_channel_first
         self.num_quantizers = num_quantizers
 
+        # soft clamping the input value
+
+        self.soft_clamp_input_value = soft_clamp_input_value
+
+        # layers
+
         self.levels = levels
         self.layers = nn.ModuleList([])
 
@@ -170,6 +177,14 @@ def forward(
 
         x = self.project_in(x)
 
+        # maybe softclamp input before residual layers
+
+        if exists(self.soft_clamp_input_value):
+            clamp_value = self.soft_clamp_input_value
+            x = (x / clamp_value).tanh() * clamp_value
+
+        # ready some variables to be accumulated
+
         quantized_out = 0.
         residual = x