src/transformers/models/pegasus/modeling_flax_pegasus.py/FlaxPegasusForConditionalGeneration/update_inputs_for_generation
class FlaxPegasusForConditionalGeneration:    def update_inputs_for_generation(self, model_outputs, model_kwargs):
        model_kwargs["past_key_values"] = model_outputs.past_key_values
        model_kwargs["decoder_position_ids"] = model_kwargs["decoder_position_ids"][:, -1:] + 1
        return model_kwargs