src/transformers/models/transfo_xl/modeling_tf_transfo_xl.py/TFTransfoXLLMHeadModel/prepare_inputs_for_generation
class TFTransfoXLLMHeadModel:    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **model_kwargs):
        inputs = {}

        # if past is defined in model kwargs then use it for faster decoding
        if past_key_values:
            input_ids = tf.expand_dims(input_ids[:, -1], axis=-1)
        else:
            input_ids = input_ids

        return inputs