Generation
Each framework has a generate method for text generation implemented in their respective GenerationMixin class:

PyTorch [~generation.GenerationMixin.generate] is implemented in [~generation.GenerationMixin].
TensorFlow [~generation.TFGenerationMixin.generate] is implemented in [~generation.TFGenerationMixin].
Flax/JAX [~generation.FlaxGenerationMixin.generate] is implemented in [~generation.FlaxGenerationMixin].