src/transformers/models/gpt_neox_japanese/modeling_gpt_neox_japanese.py/GPTNeoXJapaneseAttention/_create_causal_mask
class GPTNeoXJapaneseAttention:    def _create_causal_mask(self, key_length, query_length):
        causal_mask = torch.tril(
            torch.ones((self.max_positions, self.max_positions), dtype=torch.bool).view(
                1, 1, self.max_positions, self.max_positions
            )
        )
        return causal_mask[:, :, key_length - query_length : key_length, :key_length]