Muhan Zhong

Bitwise-Correct and Faster Top-k on TPU v4

When jax.lax.top_k is called inside a Pallas kernel on TPU, Pallas's official implementation returns duplicate indices if the input contains -inf, and a source comment states that the behavior is kept on purpose because fixing it would be too expensive. We test whether this trade...