def forward_jit(
        self,
        input: Tensor,
        lengths: Tensor,
        mems: Tensor,
        left_context_key: Tensor,
        left_context_val: Tensor,
        rpe: Optional[Tensor],
    ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
        """
        forward function for NoSegAugmentedMemoryMultiheadAttentionBmm in decoding.

        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, ...].
            left_context_key: left_context for key part. This is only used for online
                decoding. In training, this is empty tensor
            left_context_val: left_context for value part. This is only used for online
                decoding. In training, this is empty tensor

        """
        lc_length = left_context_key.size(0)

        # In decoding, summary_length = 1 or 0
        if self.use_mem:
            summary_length = 1
        else:
            summary_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

        # 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 :, :, :]

        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,
        )
        # In online decoding, we don't have attention mask. But we still need
        # to disable the attention from summary query to memory
        attention_weights_float[:, -1, :mem_length] = 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,
        )

        # In decoding, summary length is 1
        if self.use_mem:
            next_m = rc_output_memory[-1:]
            next_m = self.squash_mem(next_m)
            # rc and output
            rc_output = rc_output_memory[:-1]
            if not self.nonlinear_squash_mem:
                next_m = torch.clamp(next_m, min=-10, max=10)
        else:
            rc_output = rc_output_memory
            # empty tensor as input mems
            next_m = mems

        return rc_output, next_m, next_k, next_v

        forward function for NoSegAugmentedMemoryMultiheadAttentionBmm in decoding.

        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, ...].
            left_context_key: left_context for key part. This is only used for online
                decoding. In training, this is empty tensor
            left_context_val: left_context for value part. This is only used for online
                decoding. In training, this is empty tensor

        
structured_prompting/fairseq-version/fairseq/fairseq/models/speech_to_text/modules/emformer.py