Sparse softmax as functional?

username: AerysS
cooked: <p>I have a torch tensor of shape <code>(batch_size, N)</code>. I want to apply functional softmax with dim 1 to this tensor, but I also want it to ignore zeros in the tensor and only apply it to non-zero values (the non-zeros in the tensor are positive numbers). I think what I am looking for is the sparse softmax.</p>
<p>I came up with this code: <a href="https://github.com/KrisKorrel/sparsemax-pytorch" rel="noopener nofollow ugc">GitHub</a>, but seems like it uses <code>nn.Module</code> instead of functional.</p>
<p>How can I use that sparse softmax as a functional one (that is, it is not a layer)?</p>

username: pascal_notsawo
cooked: <p>As I see their code, something like that might do the trick.</p>
<pre><code class="lang-auto">def sparsemax(input, dim):
    """
    Args:
        input (torch.Tensor): Input tensor. First dimension should be the batch size
    Returns:
        torch.Tensor: [batch_size x number_of_logits] Output tensor
    """
    # Sparsemax currently only handles 2-dim tensors,
    # so we reshape to a convenient shape and reshape back after sparsemax
    input = input.transpose(0, dim)
    original_size = input.size()
    input = input.reshape(input.size(0), -1)
    input = input.transpose(0, 1)
    dim = 1
    ...
    ...
</code></pre>

username: AerysS
cooked: <p>So I can use the whole <code>forward()</code> function, and ignore the <code>backward()</code>, right?</p>

username: pascal_notsawo
cooked: <p>If you ignore it I think that during the back-propagation your model will calculate the gradient in the wrong way.<br>
Hence the role of backward (and thus the use of nn.Module).<br>
Let me check with an example.</p>
<p>Why don’t you want to use the version with nn.Module?</p>

username: KFrank
cooked: <p>Hi Aerys!</p>
<aside class="quote no-group" data-username="AerysS" data-post="1" data-topic="120130" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aeryss/48/37597_2.png" class="avatar"> AerysS:</div>
<blockquote>
<p>How can I use that sparse softmax as a functional one (that is, it is not a layer)?</p>
</blockquote>
</aside>
<p>I didn’t look at the github code you linked to, but, in general, if you<br>
have a (properly-written) <code>torch.nn.Module</code> or<br>
<code>torch.autograd.Function</code>, you can instantiate it “on the fly” and<br>
call it, rather than using it as a “layer” in a model.</p>
<p>Thus, for example:</p>
<pre data-code-wrap="python"><code class="lang-python">probs = torch.nn.Softmax (dim = 1) (logits)
</code></pre>
<p>is the practical equivalent of:</p>
<pre data-code-wrap="python"><code class="lang-python">probs = torch.nn.functional.softmax (logits, dim = 1)
# or
probs = torch.softmax (logits, dim = 1)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: AerysS
cooked: <p>Hi KFrank, thanks for your help. Also as pascal_notsawo said, I need it separately to apply to a tensor and do not need its gradients. So in that case, I can warp it with no_grad and it will be fine right?</p>

username: KFrank
cooked: <p>Hi Aerys!</p>
<aside class="quote no-group" data-username="AerysS" data-post="6" data-topic="120130" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aeryss/48/37597_2.png" class="avatar"> AerysS:</div>
<blockquote>
<p>I need it separately to apply to a tensor and do not need its gradients. So in that case, I can warp it with no_grad and it will be fine right?</p>
</blockquote>
</aside>
<p>Roughly speaking, whether computations are tracked for potential future<br>
gradient computation depends on the tensor being passed through the<br>
function, rather than on the function itself.</p>
<p>In short, I wouldn’t worry about it.  Even if you pass a tensor with<br>
<code>requires_grad = True</code> through your function, your function’s<br>
<code>backward()</code> will only be called if you call <code>backward()</code> on something<br>
that depends on the result of calling your function.</p>
<p>(If you do call your function on a tensor with <code>requires_grad = True</code><br>
and you know that you won’t be needing the associated gradient,<br>
you can wrap your function call in a  <code>with torch.no_grad():</code> block<br>
to save on the overhead of building that branch of the computation<br>
graph.  But I would wrap the function call rather than the innards of<br>
the function itself.)</p>
<p>Best.</p>
<p>K. Frank</p>