        """
        * `x` is the input embedding tensor $X$ of shape `[seq_len, batch_size, d_model]`
        * `mask` is a boolean mask of shape `[seq_len, seq_len, 1]` that controls the visibility of tokens
         among each other.
        """
        # Keep a copy for shortcut connection
        shortcut = x