Match floating-point comparer (checks for NaNs) in gpu_sort_rewriter pass
Match floating-point comparer (checks for NaNs) in gpu_sort_rewriter pass

This is needed to optimize the `tf.math.top_k` function with floating-point inputs.
