Get indices of the max of a 2D Tensor

username: GeoffNN
cooked: <p>I have a 2d Tensor <code>A</code> of shape <code>(d, d)</code> and I want to get the indices of its maximal element.<br>
<code>torch.argmax</code> only returns a single index. For now I got the result doing the following, but it seems involuted.</p>
<pre><code class="lang-auto">vals, row_idx = A.max(0)
col_idx = vals.argmax(0)
</code></pre>
<p>And then, <code>A[row_idx, col_idx]</code> is the correct maximal value. Is there any more straightforward way to get this?</p>

username: ptrblck
cooked: <p>Alternatively, this code should also work:</p>
<pre><code class="lang-python">x = torch.randn(10, 10)
print((x==torch.max(x)).nonzero())
</code></pre>

username: SerenityRising
cooked: <p>For anyone who stumbles in here and wonders which approach is faster (that provided by GeoffNN or ptrblck), the one-liner by ptrblck appears to be at least twice as fast.</p>
<p>In my (not rigorous) benchmarking, ptrblck’s code found the max indices of 100k random tensors in an average of 1.6 seconds, and the solution by GeoffNN found the max indices of 100k random tensors in an average of 3.5 seconds.</p>

username: GeoffNN
cooked: <p>Is there any way to batch this?</p>
<p>Suppose that now I have a 3D tensor of shape <code>(batch_size, d, d)</code> and I want to get for each 2D datapoint in the batch the index of its maximal element? I can’t figue out how to generalize the <code>(x==torch.max(x)).nonzero()</code> approach.</p>

username: David_Gilbertson
cooked: <p>Since <code>argmax()</code> gives you the index in a flattened tensor, you can infer the position in your 2D tensor from size of the last dimension. E.g. if <code>argmax()</code> returns <code>10</code> and you’ve got <code>4</code> columns, you know it’s on row 2, column 2. You can use Python’s <code>divmod</code> for this</p>
<pre data-code-wrap="py"><code class="lang-nohighlight">a = torch.randn(10, 10)
row, col = divmod(a.argmax().item(), a.shape[1])
</code></pre>
<p>This can be extended to the batch case by flattening the last two dimensions, and iterating over the results to return several 2D indexes:</p>
<pre data-code-wrap="py"><code class="lang-nohighlight">a = torch.randn(3, 2, 100)
flat_indexes = a.flatten(start_dim=-2).argmax(1)
[divmod(idx.item(), a.shape[-1]) for idx in flat_indexes]
</code></pre>

username: Jinx_L
cooked: <p>I also have the requirement to batch-wise compute the 2D cordinates of the max values of each item.<br>
A simple workaround may be:</p>
<pre><code class="lang-python"># suppose the tensor is of shape (3,2,2), 
&gt;&gt;&gt; a = torch.randn(3, 2, 2)
&gt;&gt;&gt; a
tensor([[[ 0.1450, -1.3480],
         [-0.3339, -0.5133]],

        [[ 0.6867, -0.2972],
         [ 0.8768,  0.0844]],

        [[-2.3115, -0.4549],
         [-1.5074, -0.8706]]])

# then perform batch-wise max
&gt;&gt;&gt; torch.stack([(a[i]==torch.max(a[i])).nonzero() for i in range(a.size(0))], dim=0)

tensor([[[0, 0]],

        [[1, 0]],

        [[0, 1]]])
</code></pre>

username: Inigo_Moreno
cooked: <p>In case someone wants a generic function for n dimensions:</p>
<pre data-code-wrap="python"><code class="lang-python">def torch_maxloc(x):
    d = torch.argmax(x).item()
    res = []
    for s in x.shape[::-1]:
        d, m = divmod(d, s)
        res.append(m)
    return tuple(res)[::-1]
</code></pre>