Remove the optimization to replace sort with scatter.
We will switch to allowing to specify a computation for sorting. This will make it even less likely that this optimization will kick in. TopK will specify a computation to sort into descending order, and the current logic assumes that the sorting is done into ascending order. PiperOrigin-RevId: 231557946
Loading
Please sign in to comment