Path	Lines of Code
optimum/tpu/__init__.py	3
optimum/tpu/cli.py	76
optimum/tpu/distributed_model.py	79
optimum/tpu/fsdp_v2.py	43
optimum/tpu/generation/__init__.py	2
optimum/tpu/generation/logits_process.py	34
optimum/tpu/generation/token_selector.py	116
optimum/tpu/jetstream_pt_support.py	11
optimum/tpu/model.py	43
optimum/tpu/modeling.py	47
optimum/tpu/modeling_gemma.py	877
optimum/tpu/modeling_llama.py	1124
optimum/tpu/modeling_mistral.py	1106
optimum/tpu/static_cache_xla.py	23
optimum/tpu/version.py	3
optimum/tpu/xla_logger.py	14
optimum/tpu/xla_model_parallel.py	548
optimum/tpu/xla_mp_comm.py	39