From 3f2fabcfcb0574a4d5d225bdbf2dd33e35c9886c Mon Sep 17 00:00:00 2001
From: Jan Wassenberg <janwas@google.com>
Date: Sun, 24 Mar 2024 17:13:55 -0700
Subject: [PATCH] Update todo to mention PartialSort

PiperOrigin-RevId: 618685783
---
 ops.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ops.h b/ops.h
index 7aa7b62..358509b 100644
--- a/ops.h
+++ b/ops.h
@@ -711,7 +711,7 @@ static HWY_NOINLINE HWY_MAYBE_UNUSED int SampleTopK(
     const float* HWY_RESTRICT probabilities, size_t vocab_size,
     std::mt19937& gen, float temperature, TAcceptToken& accept_token) {
   static_assert(k != 0, "");
-  // TODO(austinvhuang): Optimize this implementation.
+  // TODO: Optimize, potentially using new VQSort PartialSort.
   std::array<float, k> top_k{};  // sorted from highest [0], to lowest [k-1]
   std::array<int, k> indices{};
   for (size_t i = 0; i < vocab_size; ++i) {