Merge pull request #193 from basf/sparse_fix

fix input names in sparsemax

Author: AnFreTh

mambular/arch_utils/layer_utils/sparsemax.py

@@ -36,13 +36,13 @@ class SparsemaxFunction(Function):
     """
 
     @staticmethod
-    def forward(ctx, x, dim=-1):
+    def forward(ctx, input, dim=-1):
         """
         Forward pass of sparsemax: a normalizing, sparse transformation.
 
         Parameters
         ----------
-        x : torch.Tensor
+        input : torch.Tensor
             The input tensor on which sparsemax will be applied.
         dim : int, optional
             Dimension along which to apply sparsemax. Default is -1.
@@ -53,8 +53,8 @@ def forward(ctx, x, dim=-1):
             A tensor with the same shape as the input, with sparsemax applied.
         """
         ctx.dim = dim
-        max_val, _ = x.max(dim=dim, keepdim=True)
-        x -= max_val  # Numerical stability trick, as with softmax.
+        max_val, _ = input.max(dim=dim, keepdim=True)
+        input -= max_val  # Numerical stability trick, as with softmax.
         tau, supp_size = SparsemaxFunction._threshold_and_support(input, dim=dim)
         output = torch.clamp(input - tau, min=0)
         ctx.save_for_backward(supp_size, output)