Gumbel-(Soft)max sampling