Thank you @albanD  very much for your answer. What if i take a weighted sum of all the elements in matrix (softmax) but then multiplied the those weights with a mask generated based on the idx, as follows:
weights = softmax(A.W.B)
_, idx = torch.max(weights)
mask = torch.ones_like(weights)
mask = mask[weights<idx] = 0
weighted_sum = torch.sum(weights * mask * matrix)

This would basically achieve the same purpose as before, because ignored elements are multiplied with 0. So Would that be differentiable?