Change tensor values by index lower than some value

isg: <p>Hi all,</p>
<p>I’m new and have a simple problem I cannot solve. I have a (2,n) tensor of logits and I would like to do a topk selection. I’ve been reading <a href="https://www.packtpub.com/en-us/learning/how-to-tutorials/exploring-token-generation-strategies?srsltid=AfmBOoo_z6FRUAwDJnZ2MbChTXW7SN67Ua3OtLYGyM6RUJb8BayQNFUq" rel="noopener nofollow ugc">this link</a> to implement what I want. The code works when the logits are a 1d but not when they’re (2,n).<br>
I can’t figure out how to modify the code accordingly.</p>
<p>There is a very similar post on this forum using the cumsum but is not what I want. My minimal code is:</p>
<pre><code class="lang-auto">logits = torch.tensor([[ 0.0333,  0.5174,  0.4797],[-0.0644, -0.5815, -0.7552]])
print(logits)
topk_values, topk_indices = torch.topk(logits, 2)
last = topk_values[:,-1].unsqueeze(dim=-1)
print(last)
print(topk_indices)
logits[logits &lt; topk_values[-1]] = float('-inf')
</code></pre>
<p>but I’m getting the following error:</p>
<pre><code class="lang-auto">RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 1
</code></pre>
<p>Could you please help me to understand what’s going on and how to solve it?</p>
<p>Thank you!</p>

KFrank: <p>Hi isg!</p>
<aside class="quote no-group" data-username="isg" data-post="1" data-topic="217421">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/i/53a042/48.png" class="avatar"> isg:</div>
<blockquote>
<pre><code class="lang-auto">logits[logits &lt; topk_values[-1]] = float('-inf')
</code></pre>
<p>but I’m getting the following error:</p>
<pre><code class="lang-auto">RuntimeError: The size of tensor a (3) must match the size of tensor b (2) at non-singleton dimension 1
</code></pre>
</blockquote>
</aside>
<p>The problem is that <code>logits</code>, with trailing dimension 3, and <code>topk_values[-1]]</code>, with<br>
trailing dimension 2, are not <a href="https://pytorch.org/docs/2.6/notes/broadcasting.html" rel="noopener nofollow ugc">broadcastable.</a></p>
<p>I’m guessing what you want is:</p>
<pre data-code-wrap="python"><code class="lang-python">logits[logits &lt; last] = float('-inf')
</code></pre>
<p>In this case, <code>unsqueeze()</code> has been used to add a trailing singleton (size = 1) dimension<br>
to <code>last</code>, which is broadcastable to the size-3 trailing dimension of <code>logits</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

isg: <p>Thank you <a class="mention" href="/u/kfrank">@KFrank</a>!</p>
<p>Correctly, my problem was assuming some broadcasting that was impossible. You pointed me in the right direction (without giving me the solution, which forced me to read more about broadcasting and learn more, thanks again).</p>
<p>At the end, the solution that I have implemented is a class to do what I wanted.</p>
<pre><code class="lang-auto">class TopK:

    def __init__(self, k: int) -&gt; None:
        
        self.k = k

    def predict(self, logits) -&gt; torch.tensor:

        # We reshape the logits tensor removing the leading dimension: "batch"
        logits = logits.view(logits.shape[0]*logits.shape[1],-1)

        # We compute the topk values
        topk_values, topk_indices = torch.topk(logits, self.k)

        # We compute the minimum value for each row
        min_values = torch.min(topk_values, dim=-1).values.view(logits.shape[0],1)

        # We renormalize the probabilities setting any logits smaller than
        # the topk to "-inf" first, and then computing the softmax row-wise
        logits[logits &lt; min_values] = float('-inf')
        probs = logits.softmax(dim=-1)

        # Finally, we use the multinomial Pytorch method to make a random pick across the
        # topk using the re-normalized probabilities.
        results = torch.multinomial(probs, num_samples=1).flatten()

        return results
</code></pre>