Hi,
Gradients will flow back from desired_element to matrix as you only take one element out of a matrix which is differentiable.
Gradients won’t flow back towards softmax though as indexing is not differentiable wrt the index and the argmax operation is not differentiable either.