How does torch.argmax work for 4-dimensions
I am a newbie in Pytorch. Even though I read the documentation, it is unclear for me how does torch.argmax() applied to first dimension work when we have 4-dimensional input. Also, how does keepdims=True change the output?
Here is an example of each case:
k = torch.rand(2, 3, 4, 4)
print(k):

tensor([[[[0.2912, 0.4818, 0.1123, 0.3196],
          [0.6606, 0.1547, 0.0368, 0.9475],
          [0.4753, 0.7428, 0.5931, 0.3615],
          [0.6729, 0.7069, 0.1569, 0.3086]],

         [[0.6603, 0.7777, 0.3546, 0.2850],
          [0.3681, 0.5295, 0.8812, 0.6093],
          [0.9165, 0.2842, 0.0260, 0.1768],
          [0.9371, 0.9889, 0.6936, 0.7018]],

         [[0.5880, 0.0349, 0.0419, 0.3913],
          [0.5884, 0.9408, 0.1707, 0.1893],
          [0.3260, 0.4410, 0.6369, 0.7331],
          [0.9448, 0.7130, 0.3914, 0.2775]]],


        [[[0.9433, 0.8610, 0.9936, 0.1314],
          [0.8627, 0.3103, 0.3066, 0.3547],
          [0.3396, 0.1892, 0.0385, 0.5542],
          [0.4943, 0.0256, 0.7875, 0.5562]],

         [[0.2338, 0.2498, 0.4749, 0.2520],
          [0.4405, 0.1605, 0.6219, 0.8955],
          [0.2326, 0.1816, 0.5032, 0.8732],
          [0.2089, 0.6131, 0.1898, 0.0517]],

         [[0.1472, 0.8059, 0.6958, 0.9047],
          [0.6403, 0.2875, 0.5746, 0.5908],
          [0.8668, 0.4602, 0.8224, 0.9307],
          [0.2077, 0.5665, 0.8671, 0.4365]]]])

argmax = torch.argmax(k, axis=1)
print(argmax):
tensor([[[1, 1, 1, 2],
         [0, 2, 1, 0],
         [1, 0, 2, 2],
         [2, 1, 1, 1]],

        [[0, 0, 0, 2],
         [0, 0, 1, 1],
         [2, 2, 2, 2],
         [0, 1, 2, 0]]])


argmax = torch.argmax(k, axis=1, keepdims=True)
print(argmax):
tensor([[[[1, 1, 1, 2],
          [0, 2, 1, 0],
          [1, 0, 2, 2],
          [2, 1, 1, 1]]],


        [[[0, 0, 0, 2],
          [0, 0, 1, 1],
          [2, 2, 2, 2],
          [0, 1, 2, 0]]]])


If k is a tensor of shape (2, 3, 4, 4), by definition, torch.argmax with axis=1 should give you an output of shape (2, 4, 4). To understand why this happens, you have to understand what happens in lower dimensions first.
If I have a 2D (2, 2) tensor A, like:
[[1,2],
 [3,4]]

Then torch.argmax(A, axis=1) gives the output of shape (2) with values (1, 1). The axis argument means axis along which to operate. So setting axis=1 means that it will look at values from each column one by one, before deciding a max. For row 0, it looks at column values 1, 2 and decides that 2 (at index 1) is the max. For row 1, it looks at column vales 3, 4 and decides that 4 (at index 1) is the max. So the argmax result is [1, 1].
Moving up to 3D, let's have a hypothetical array of dimensions (I, J, K). If we call argmax with axis = 1, we can break it down to the following:
I, J, K = 3, 4, 5
A = torch.rand(I, J, K)
out = torch.zeros((I, K), dtype=torch.int32)

for i in range(I):
    for k in range(K):
        out[i,k] = torch.argmax(A[i,:,k])
        
print(out)
print(torch.argmax(A, axis=1))

Out:
tensor([[3, 3, 2, 3, 2],
        [1, 1, 0, 1, 0],
        [0, 1, 0, 3, 3]], dtype=torch.int32)
tensor([[3, 3, 2, 3, 2],
        [1, 1, 0, 1, 0],
        [0, 1, 0, 3, 3]])

So what happens is, in your 3D tensor, you're once again calculating argmax along the columns/axis 1. So for each unique pair of (i, k), you have exactly J values along the axis 1, right? The index of the maximum value within those J values is inserted into position (i,k) of the output.
If you understand this, then you can understand what happens in 4D. For any 4D tensor of dimensions (I, J, K, L), if you call argmax with axis=1, then for each combination of (i, k, l) you'll have exactly J values along axis 1 - and the argmax of those J values will be present at output[i,k,l].
The keepdims argument is merely conserving the number of dimensions of your matrix. For example, argmax at axis 1 on the 4D matrix gives a 3D result of shape (I,K,L), but using keepdims, the result will be 4D as well with the shape (I,1,K,L).
