Merge pull request #194 from basf/develop

Update feature branch

Author: mkumar73

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)
@@ -86,7 +86,7 @@ def backward(ctx, grad_output):  # type: ignore
         return grad_input, None
 
     @staticmethod
-    def _threshold_and_support(x, dim=-1):
+    def _threshold_and_support(input, dim=-1):
         """
         Computes the threshold and support for sparsemax.
 
@@ -103,14 +103,14 @@ def _threshold_and_support(x, dim=-1):
             - torch.Tensor : The threshold value for sparsemax.
             - torch.Tensor : The support size tensor.
         """
-        input_srt, _ = torch.sort(x, descending=True, dim=dim)
+        input_srt, _ = torch.sort(input, descending=True, dim=dim)
         input_cumsum = input_srt.cumsum(dim) - 1
         rhos = _make_ix_like(input, dim)
         support = rhos * input_srt > input_cumsum
 
         support_size = support.sum(dim=dim).unsqueeze(dim)
         tau = input_cumsum.gather(dim, support_size - 1)
-        tau /= support_size.to(x.dtype)
+        tau /= support_size.to(input.dtype)
         return tau, support_size