optimum/tpu/generation/token_selector.py
text-generation-inference/server/text_generation_server/jetstream_pt_support/token_selector.py

optimum/tpu/generation/logits_process.py
text-generation-inference/server/text_generation_server/jetstream_pt_support/logits_process.py