def forward_jit(
        self, input: Tensor, lengths: Tensor, state: Optional[List[Tensor]] = None
    ) -> Tuple[Tensor, Tensor, List[Tensor]]:
        """
        Forward helper for online decoding.

        args:
            input: [seg, right_context]. We assume in online we
                always padding the right context to the preset right context size.
                For the last segment, we may have short segment size, but right
                context size is the same as other segments
            lengths: utterance input length is the utterance segment length and
                     right context size
            state: [memory, left_context_key, left_context_val]. To improve throughput,
                in addition to memory, we also cache key and value for left_context in
                multihead self-attention
        """
        # In online decoding, input = [segment, right_context]
        # Lengths = [segment_length, right_context_length]
        # so we need strip right context in output
        T, B, D = input.size()
        rc_str = T - self.right_context
        rc_end = T
        right_context_blocks = input[rc_str:rc_end, :, :]
        seg = input[:rc_str, :, :]
        lengths = torch.clamp(lengths - self.right_context, min=0)
        mems_list = []
        lc_key_list = []
        lc_val_list = []

        results = self.forward_jit_mini_batch_init(seg, state, True)
        state, mems, state_mems, past_length, past_left_context = results

        # relative position embedding
        if self.use_rpe:
            rpe = self._get_relative_position(
                input=input,
                max_relative_position=self.max_relative_position,
                left_context_length=past_left_context,
                past_length=past_length,
                is_decoding=True,
            )
        else:
            rpe = None

        # memory for first layer.
        mems_list.append(state_mems[-self.max_memory_size :, :, :])
        output = seg
        i = 0
        for layer in self.layers:
            # In order to make cross stream batching work, mem, left context key
            # and left context value in the state should always be the same shape.
            # We use the past length to track the processed segment number. In this
            # way, we take out the essential memory, left context key and left
            # context val from the state. After finish the forward for current segment
            # we add the new memory, left context key and left context value into the
            # staate and trim out the oldest part to keep the shape consistent.
            true_mems, lc_key, lc_val = self.state_update_before(
                layer=i,
                state=state,
                past_length=past_length,
                past_left_context=past_left_context,
            )

            output, mems, right_context_blocks, next_key, next_val = layer.forward_jit(
                input=output,
                lengths=lengths,
                mems=true_mems,
                right_context_blocks=right_context_blocks,
                left_context_key=lc_key,
                left_context_val=lc_val,
                rpe=rpe,
            )
            # mems is used for next layer
            mems_list, lc_key_list, lc_val_list, _ = self.state_update_after(
                layer=i,
                state=state,
                mems_list=mems_list,
                mems=mems,
                next_key=next_key,
                next_val=next_val,
                lc_key_list=lc_key_list,
                lc_val_list=lc_val_list,
            )
            i += 1

        # update state
        state = self.state_update_after_loop(
            state=state,
            mems_list=mems_list,
            lc_key_list=lc_key_list,
            lc_val_list=lc_val_list,
            update_length=1,
        )

        return output, lengths, state

        Forward helper for online decoding.

        args:
            input: [seg, right_context]. We assume in online we
                always padding the right context to the preset right context size.
                For the last segment, we may have short segment size, but right
                context size is the same as other segments
            lengths: utterance input length is the utterance segment length and
                     right context size
            state: [memory, left_context_key, left_context_val]. To improve throughput,
                in addition to memory, we also cache key and value for left_context in
                multihead self-attention
        
structured_prompting/fairseq-version/fairseq/fairseq/models/speech_to_text/modules/emformer.py