def _update_attention(model_kwargs, new_past_index, is_encoder_decoder):
            """updates the appropriate attention mask -- encoder-decoder models use `decoder_attention_mask`"""
            update_start = tf.constant([0, 1], dtype=tf.int32) * new_past_index
            if is_encoder_decoder:
                decoder_attention_mask = model_kwargs.pop("decoder_attention_mask")
                decoder_attention_mask_update_slice = tf.ones((batch_size, 1), dtype=decoder_attention_mask.dtype)
                decoder_attention_mask = dynamic_update_slice(
                    decoder_attention_mask, decoder_attention_mask_update_slice, update_start
                )
                mask = {"decoder_attention_mask": decoder_attention_mask}
            else:
                attention_mask = model_kwargs.pop("attention_mask")
                attention_mask_update_slice = tf.ones((batch_size, 1), dtype=attention_mask.dtype)
                attention_mask = dynamic_update_slice(attention_mask, attention_mask_update_slice, update_start)
                mask = {"attention_mask": attention_mask}
            return mask
updates the appropriate attention mask -- encoder-decoder models use `decoder_attention_mask`
dpkd/transformers/src/transformers/generation/tf_utils.py