    def forward(
        self,
        input: Tensor,
        lengths: Tensor,
        mems: Tensor,
        attention_mask: Tensor,
        pre_mems: Optional[Tensor] = None,
        left_context_key: Optional[Tensor] = None,
        left_context_val: Optional[Tensor] = None,
        rpe: Optional[Tensor] = None,
    ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
        """
        forward function for NoSegAugmentedMemoryMultiheadAttentionBmm in training.

        args:
            input: formed in the following way
                [right_context_0, right_contex_1, ..., seg_0, seg_1,
                ..., summary_0, summary_1,..]
            lengths: the length of query which is [seg_0, seg_1, ....]
            mems: [mem_0, mem_1, ...].
            attention_mask: attention mask for query = [right_context, query, summary]
                key = [mem, right_context, query]. This is only used for traing.

        """
        if self.use_mem:
            mem_length = mems.size(0)
            summary_length = mem_length + 1
            if pre_mems is not None:
                mems = torch.cat([pre_mems, mems], dim=0)
        else:
            mem_length = 0
            summary_length = 0

        # In training, lc_length = 0
        if left_context_key is not None:
            lc_length = left_context_key.size(0)
        else:
            lc_length = 0
        results = self.prepare_qkv(
            input=input,
            mems=mems,
            lengths=lengths,
            summary_length=summary_length,
            lc_length=lc_length,
        )
        result_qkv, input_shape, result_lengths_info, padding_mask = results
        q, k, v = result_qkv
        (
            mem_length,
            utterance_length,
            right_context_blocks_length,
            key_length,
        ) = result_lengths_info

        if left_context_key is not None:
            # add the cache key and value
            new_k = torch.cat(
                [
                    k[: mem_length + right_context_blocks_length, :, :],
                    left_context_key,
                    k[-utterance_length:, :, :],
                ],
                dim=0,
            )
            new_v = torch.cat(
                [
                    v[: mem_length + right_context_blocks_length, :, :],
                    left_context_val,
                    v[-utterance_length:, :, :],
                ],
                dim=0,
            )
            next_k = new_k[mem_length + right_context_blocks_length :, :, :]
            next_v = new_v[mem_length + right_context_blocks_length :, :, :]
        else:
            new_k = k
            new_v = v
            next_k = None
            next_v = None

        attention_weights, attention_weights_float, v = self.prepare_attention_weights(
            q=q,
            new_k=new_k,
            new_v=new_v,
            input_shape=input_shape,
            rpe=rpe,
        )

        # mask attention
        attention_mask = attention_mask.unsqueeze(0)
        attention_weights_float = attention_weights_float.masked_fill(
            attention_mask, float(self.negative_inf)
        )

        rc_output_memory = self.prepare_attention_output(
            attention_weights=attention_weights,
            attention_weights_float=attention_weights_float,
            v=v,
            input_shape=input_shape,
            key_length=key_length,
            padding_mask=padding_mask,
            rpe=rpe,
        )

        if self.use_mem:
            # next_m length equals to summary length - 1
            # last memory is ignored
            if self.mini_batches:
                next_m = rc_output_memory[-summary_length:]
            else:
                next_m = rc_output_memory[-summary_length:-1]

            next_m = self.squash_mem(next_m)
            # rc and output
            rc_output = rc_output_memory[:-summary_length]
            if not self.nonlinear_squash_mem:
                next_m = torch.clamp(next_m, min=-10, max=10)
        else:
            next_m = mems
            rc_output = rc_output_memory

        return rc_output, next_m, next_k, next_v
forward function for NoSegAugmentedMemoryMultiheadAttentionBmm in training.

args:
input: formed in the following way
[right_context_0, right_contex_1, ..., seg_0, seg_1,
..., summary_0, summary_1,..]
lengths: the length of query which is [seg_0, seg_1, ....]
mems: [mem_0, mem_1, ...].
attention_mask: attention mask for query = [right_context, query, summary]
key = [mem, right_context, query]. This is only used for traing.
fairseq/models/speech_to_text/modules/emformer.py