src/transformers/models/gptj/modeling_flax_gptj.py/FlaxGPTJForCausalLM/update_inputs_for_generation
class FlaxGPTJForCausalLM:    def update_inputs_for_generation(self, model_outputs, model_kwargs):
        model_kwargs["past_key_values"] = model_outputs.past_key_values
        model_kwargs["position_ids"] = model_kwargs["position_ids"][:, -1:] + 1
        return model_kwargs