feat: pass step lut to identity step function and improve readablility

elasticai/creator/base_modules/tanh.py

@@ -17,16 +17,15 @@ def __init__(
     ) -> None:
         super().__init__()
         self._arithmetics = arithmetics
-        self.num_steps = num_steps
-        self.sampling_intervall = sampling_intervall
+        self._step_lut = torch.linspace(*sampling_intervall, num_steps)
 
-    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+    def _quantized_step_inputs(self, inputs: torch.Tensor) -> torch.Tensor:
         step_inputs = cast(
-            torch.Tensor,
-            IdentityStepFunction.apply(
-                inputs, *self.sampling_intervall, self.num_steps
-            ),
+            torch.Tensor, IdentityStepFunction.apply(inputs, self._step_lut)
         )
-        quantized_step_inputs = self._arithmetics.quantize(step_inputs)
-        outputs = torch.nn.functional.tanh(quantized_step_inputs)
+        return self._arithmetics.quantize(step_inputs)
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        inputs = self._quantized_step_inputs(inputs)
+        outputs = torch.nn.functional.tanh(inputs)
         return self._arithmetics.quantize(outputs)