Make TopK stable on the GPU backend.

TopK op needs the guarantee that the sorting is stable. In that case, a
'values' parameter is passed which is a iota. We can use that to disambiguate
comparison ties to sort by original index, which makes the sorting stable.
Note that this still doesn't guarantee that the sort op is stable for all
inputs, just that it is stable when used by TopK.

PiperOrigin-RevId: 222390450