Differentiable argmax

What is the trick you have for this? I believe @ram’s problem is similar to the one I have here: Predict a categorical variable and then embed it (one-hot?)