How to go from 1d index that results from flatten operation to 2d index?

1

Say we have a matrix M that is n by n and after we flatten it we get n^2 sized array.
Then we find max at index k.
How to convert k to (x, y) coordinate for pre-flatten array?
Is there torch function for this?

2 
n = 5
M = torch.rand(n, n)
flatten_M = M.view(-1)
k = torch.argmax(flatten_M)
print(k, flatten_M[k])
i, j = divmod(k.item(), n)
print(M[i, j])