    def forward_chunk(
        self,
        xs: torch.Tensor,
        offset: int,
        required_cache_size: int,
        att_cache: torch.Tensor = torch.zeros(0, 0, 0, 0),
        cnn_cache: torch.Tensor = torch.zeros(0, 0, 0, 0),
        att_mask: torch.Tensor = torch.ones((0, 0, 0), dtype=torch.bool),
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Forward just one chunk

        Args:
            xs (torch.Tensor): chunk input
            offset (int): current offset in encoder output time stamp
            required_cache_size (int): cache size required for next chunk
                compuation
                >=0: actual cache size
                <0: means all history cache is required
            att_cache (torch.Tensor): cache tensor for KEY & VALUE in
                transformer/conformer attention, with shape
                (elayers, head, cache_t1, d_k * 2), where
                `head * d_k == hidden-dim` and
                `cache_t1 == chunk_size * num_decoding_left_chunks`.
            cnn_cache (torch.Tensor): cache tensor for cnn_module in conformer,
                (elayers, b=1, hidden-dim, cache_t2), where
                `cache_t2 == cnn.lorder - 1`
            att_mask : mask matrix of self attention

        Returns:
            torch.Tensor: output of current input xs
            torch.Tensor: subsampling cache required for next chunk computation
            List[torch.Tensor]: encoder layers output cache required for next
                chunk computation
            List[torch.Tensor]: conformer cnn cache

        """
        assert xs.size(0) == 1

        # using downsampling factor to recover offset
        offset *= self.calculate_downsampling_factor(self.num_blocks + 1)

        chunk_masks = torch.ones(1, xs.size(1), device=xs.device, dtype=torch.bool)
        chunk_masks = chunk_masks.unsqueeze(1)  # (1, 1, xs-time)

        real_len = 0
        if self.global_chunk_size > 0:
            # for ONNX decode simulation， padding xs to chunk_size
            real_len = xs.size(1)
            pad_len = self.chunk_feature_map - real_len
            xs = F.pad(xs, (0, 0, 0, pad_len), value=0.0)
            chunk_masks = F.pad(chunk_masks, (0, pad_len), value=0.0)

        if self.global_cmvn is not None:
            xs = self.global_cmvn(xs)

        # NOTE(xcsong): Before embed, shape(xs) is (b=1, time, mel-dim)
        xs, pos_emb, chunk_masks = self.embed(xs, chunk_masks, offset)
        elayers, cache_t1 = att_cache.size(0), att_cache.size(2)
        chunk_size = xs.size(1)
        attention_key_size = cache_t1 + chunk_size
        # NOTE(xcsong): After  embed, shape(xs) is (b=1, chunk_size, hidden-dim)
        # shape(pos_emb) = (b=1, chunk_size, emb_size=output_size=hidden-dim)

        if required_cache_size < 0:
            next_cache_start = 0
        elif required_cache_size == 0:
            next_cache_start = attention_key_size
        else:
            next_cache_start = max(attention_key_size - required_cache_size, 0)

        r_att_cache = []
        r_cnn_cache = []
        mask_pad = torch.ones(1, xs.size(1), device=xs.device, dtype=torch.bool)
        mask_pad = mask_pad.unsqueeze(1)  # batchPad (b=1, 1, time=chunk_size)

        if self.global_chunk_size > 0:
            # for ONNX decode simulation
            pos_emb = self.embed.position_encoding(
                offset=max(offset - cache_t1, 0), size=cache_t1 + self.global_chunk_size
            )
            att_mask[:, :, -self.global_chunk_size :] = chunk_masks
            mask_pad = chunk_masks.to(torch.bool)
        else:
            pos_emb = self.embed.position_encoding(
                offset=offset - cache_t1, size=attention_key_size
            )

        max_att_len, max_cnn_len = 0, 0  # for repeat_interleave of new_att_cache
        for i, layer in enumerate(self.encoders):
            factor = self.calculate_downsampling_factor(i)
            # NOTE(xcsong): Before layer.forward
            #   shape(att_cache[i:i + 1]) is (1, head, cache_t1, d_k * 2),
            #   shape(cnn_cache[i])       is (b=1, hidden-dim, cache_t2)
            # shape(new_att_cache) = [ batch, head, time2, outdim//head * 2 ]
            att_cache_trunc = 0
            if xs.size(1) + att_cache.size(2) / factor > pos_emb.size(1):
                # The time step is not divisible by the downsampling multiple
                att_cache_trunc = (
                    xs.size(1) + att_cache.size(2) // factor - pos_emb.size(1) + 1
                )
            xs, _, new_att_cache, new_cnn_cache = layer(
                xs,
                att_mask,
                pos_emb,
                mask_pad=mask_pad,
                att_cache=att_cache[i : i + 1, :, ::factor, :][
                    :, :, att_cache_trunc:, :
                ],
                cnn_cache=cnn_cache[i, :, :, :] if cnn_cache.size(0) > 0 else cnn_cache,
            )

            if i in self.stride_layer_idx:
                # compute time dimension for next block
                efficient_index = self.stride_layer_idx.index(i)
                att_mask = att_mask[
                    :, :: self.stride[efficient_index], :: self.stride[efficient_index]
                ]
                mask_pad = mask_pad[
                    :, :: self.stride[efficient_index], :: self.stride[efficient_index]
                ]
                pos_emb = pos_emb[:, :: self.stride[efficient_index], :]

            # shape(new_att_cache) = [batch, head, time2, outdim]
            new_att_cache = new_att_cache[:, :, next_cache_start // factor :, :]
            # shape(new_cnn_cache) = [1, batch, outdim, cache_t2]
            new_cnn_cache = new_cnn_cache.unsqueeze(0)

            # use repeat_interleave to new_att_cache
            new_att_cache = new_att_cache.repeat_interleave(repeats=factor, dim=2)
            # padding new_cnn_cache to cnn.lorder for casual convolution
            new_cnn_cache = F.pad(
                new_cnn_cache, (self.cnn_module_kernel - 1 - new_cnn_cache.size(3), 0)
            )

            if i == 0:
                # record length for the first block as max length
                max_att_len = new_att_cache.size(2)
                max_cnn_len = new_cnn_cache.size(3)

            # update real shape of att_cache and cnn_cache
            r_att_cache.append(new_att_cache[:, :, -max_att_len:, :])
            r_cnn_cache.append(new_cnn_cache[:, :, :, -max_cnn_len:])

        if self.normalize_before:
            xs = self.after_norm(xs)

        # NOTE(xcsong): shape(r_att_cache) is (elayers, head, ?, d_k * 2),
        #   ? may be larger than cache_t1, it depends on required_cache_size
        r_att_cache = torch.cat(r_att_cache, dim=0)
        # NOTE(xcsong): shape(r_cnn_cache) is (e, b=1, hidden-dim, cache_t2)
        r_cnn_cache = torch.cat(r_cnn_cache, dim=0)

        if self.global_chunk_size > 0 and real_len:
            chunk_real_len = (
                real_len
                // self.embed.subsampling_rate
                // self.calculate_downsampling_factor(self.num_blocks + 1)
            )
            # Keeping 1 more timestep can mitigate information leakage
            #   from the encoder caused by the padding
            xs = xs[:, : chunk_real_len + 1, :]

        return xs, r_att_cache, r_cnn_cache
Forward just one chunk

Args:
xs (torch.Tensor): chunk input
offset (int): current offset in encoder output time stamp
required_cache_size (int): cache size required for next chunk
compuation
>=0: actual cache size
<0: means all history cache is required
att_cache (torch.Tensor): cache tensor for KEY & VALUE in
transformer/conformer attention, with shape
(elayers, head, cache_t1, d_k * 2), where
`head * d_k == hidden-dim` and
`cache_t1 == chunk_size * num_decoding_left_chunks`.
cnn_cache (torch.Tensor): cache tensor for cnn_module in conformer,
(elayers, b=1, hidden-dim, cache_t2), where
`cache_t2 == cnn.lorder - 1`
att_mask : mask matrix of self attention

Returns:
torch.Tensor: output of current input xs
torch.Tensor: subsampling cache required for next chunk computation
List[torch.Tensor]: encoder layers output cache required for next
chunk computation
List[torch.Tensor]: conformer cnn cache
modules/wenet_extractor/efficient_conformer/encoder.py