src/transformers/models/gptj/modeling_tf_gptj.py/TFGPTJAttention/get_causal_mask
class TFGPTJAttention:    def get_causal_mask(self, key_length, query_length) -> tf.Tensor:
        return tf.cast(self.lower_triangle_mask[:, :, key_length - query_length : key_length, :key_length], tf.bool)