happyz
/
gemma.cpp
mirror of https://github.com/google/gemma.cpp.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Update todo to mention PartialSort 

PiperOrigin-RevId: 618685783
This commit is contained in:
Jan Wassenberg 2024-03-24 17:13:55 -07:00 committed by Copybara-Service
parent fcf5c1af88
commit 3f2fabcfcb
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
ops.h
View File

@ -711,7 +711,7 @@ static HWY_NOINLINE HWY_MAYBE_UNUSED int SampleTopK(
const float* HWY_RESTRICT probabilities, size_t vocab_size, const float* HWY_RESTRICT probabilities, size_t vocab_size,
std::mt19937& gen, float temperature, TAcceptToken& accept_token) { std::mt19937& gen, float temperature, TAcceptToken& accept_token) {
static_assert(k != 0, ""); 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<float, k> top_k{}; // sorted from highest [0], to lowest [k-1]
std::array<int, k> indices{}; std::array<int, k> indices{};
for (size_t i = 0; i < vocab_size; ++i) { for (size_t i = 0; i < vocab_size; ++i) {