fairseq/models/speech_to_text/modules/emformer.py

#!/usr/bin/env python3
# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the LICENSE file in
# the root directory of this source tree. An additional grant of patent rights
# can be found in the PATENTS file in the same directory.


import math
import re
from functools import partial
from typing import List, Optional, Tuple

import torch
import torch.nn as nn
from fairseq.models import (
    FairseqEncoder,
)
from fairseq.models.speech_to_text.utils import (
    NoOp,
    lengths_to_padding_mask,
    segments_to_sequence,
)
from fairseq.models.speech_to_text.utils import (
    attention_suppression,
    layer_norm_backward_hook,
)
from torch import Tensor, device as Device
from torch.quantization.qconfig import (
    default_dynamic_qconfig,
    per_channel_dynamic_qconfig,
)


class RelativePositionEmbedding(nn.Module):
    """
    Implementation according to https://arxiv.org/abs/1803.02155
    """

    def __init__(self, head_dim, max_position, norm_init=True):
        super().__init__()
        self.head_dim = head_dim
        self.max_position = max_position
        self.embeddings = nn.Parameter(torch.Tensor(max_position * 2 + 1, head_dim))
        if norm_init:
            nn.init.xavier_normal_(self.embeddings)
        else:
            nn.init.xavier_uniform_(self.embeddings)

    def forward(self, input: Tensor):
        output = nn.functional.embedding(input.long(), self.embeddings)
        return output


class Fp32LayerNorm(nn.Module):
    def __init__(
        self,
        input_dim,
        clamp_grad=True,
        max_grad_value=256,
        eps=1e-5,
        elementwise_affine=True,
    ):
        super().__init__()
        self.torch_module = torch.nn.LayerNorm(
            input_dim, eps=eps, elementwise_affine=elementwise_affine
        )
        if clamp_grad:
            hook = partial(layer_norm_backward_hook, clamp_value=max_grad_value)
            self.torch_module.register_backward_hook(hook)

    def forward(self, input):
        output = torch.nn.functional.layer_norm(
            input.float(),
            self.torch_module.normalized_shape,
            self.torch_module.weight.float()
            if self.torch_module.weight is not None
            else None,
            self.torch_module.bias.float()
            if self.torch_module.bias is not None
            else None,
            self.torch_module.eps,
        ).type_as(input)
        return output


# ------------------------------------------------------------------------------
#   PositionwiseFF
# ------------------------------------------------------------------------------


class PositionwiseFF(nn.Module):
    """
    FFN layer in transformer.

    Args:
        input_dim: input embedding dimension
        ffn_dim: FFN layer inner dimension
        dropout_on_fc1: dropout for first linear layer
        dropout_on_fc2: dropout fr second linear layer
        activation_fn: activation function used after first linear layer. \
                Only relu or gelu is supported.

    """

    def __init__(
        self, input_dim, ffn_dim, dropout_on_fc1, dropout_on_fc2, activation_fn
    ):
        super(PositionwiseFF, self).__init__()

        self.input_dim = input_dim
        self.ffn_dim = ffn_dim
        if activation_fn == "relu":
            ac = nn.ReLU()
        elif activation_fn == "gelu":
            ac = nn.GELU()
        else:
            raise ValueError("Unsupported activation_fn = ({})".format(activation_fn))

        # fc1 -> ac -> dropout -> fc2 -> dropout
        self.module = nn.Sequential(
            nn.Linear(input_dim, ffn_dim),
            ac,
            nn.Dropout(dropout_on_fc1),
            nn.Linear(ffn_dim, input_dim),
            nn.Dropout(dropout_on_fc2),
        )

        self.layer_norm = Fp32LayerNorm(input_dim)

    def forward(self, input):
        module_out = self.module(self.layer_norm(input))
        output = module_out + input

        return output

    def quantize_(self, params=None):
        if params and "per_channel" in params and params["per_channel"]:
            qconfig = per_channel_dynamic_qconfig
        else:
            qconfig = default_dynamic_qconfig
        torch.quantization.quantize_dynamic(
            self, {torch.nn.Linear: qconfig}, dtype=torch.qint8, inplace=True
        )
        return self


# ------------------------------------------------------------------------------
#   SummarizationLayer
# ------------------------------------------------------------------------------


class SummarizationLayer(nn.Module):
    def __init__(self, method, segment_size, embedding_dim):
        super(SummarizationLayer, self).__init__()
        self.segment_size = segment_size
        self.embedding_dim = embedding_dim
        nonlin_match = re.match(r"nonlinear\((?P<act>[a-z]+),(?P<dim>[0-9]+)\)", method)
        self.method = method
        if method == "mean":
            self.module = nn.AvgPool1d(
                kernel_size=segment_size,
                stride=segment_size,
                ceil_mode=True,
            )
        elif method == "max":
            self.module = nn.MaxPool1d(
                kernel_size=segment_size,
                stride=segment_size,
                ceil_mode=True,
            )
        elif method == "linear":
            self.module = nn.Linear(segment_size, 1)
        elif nonlin_match:
            nonlin_args = nonlin_match.groupdict()
            act_type = nonlin_args["act"]
            hid_dim = int(nonlin_args["dim"])
            if act_type == "relu":
                act = nn.ReLU()
            elif act_type == "gelu":
                act = nn.GELU()
            else:
                raise ValueError("Unsupported activation_fn = ({})".format(act_type))
            self.module = nn.Sequential(
                nn.Linear(segment_size, hid_dim),
                act,
                nn.Linear(hid_dim, 1),
            )
        else:
            raise ValueError("Unsupported summarization method = ({})".format(method))

    def forward(self, input):
        # T, B, D -> B, D, T
        input = input.permute(1, 2, 0)

        if self.method == "mean" or self.method == "max":
            output = self.module(input)
            output = output.permute(2, 0, 1)
            return output

        full_seg_length = input.size(2) // self.segment_size * self.segment_size
        if full_seg_length > 0:
            # at least one seg is full
            B = input.size(0)
            D = input.size(1)
            input_todo = (
                input[:, :, :full_seg_length]
                .contiguous()
                .view(B, -1, self.segment_size)
            )
            output = self.module(input_todo)
            output = output.view(B, D, -1)
        else:
            output = input.new_zeros(input.size(0), input.size(1), 0)
        left = input.size(2) - full_seg_length
        if left > 0:
            # when last seg is not full, use zeros as last memory placeholder
            zeros = input.new_zeros(input.size(0), input.size(1), 1)
            output = torch.cat([output, zeros], dim=2)
        output = output.permute(2, 0, 1)
        return output


# ------------------------------------------------------------------------------
#   NoSegAugmentedMemoryMultiheadAttentionBmm
# ------------------------------------------------------------------------------


class NoSegAugmentedMemoryMultiheadAttentionBmm(nn.Module):
    """
    Whole utterance augmented memory multihead attention using BMM.

    Different with previous augmented memory multihead attention where
    the utterance is chunked into segments. Here we use attention mask
    achieve so. The input embedding [right_context, utterance, summary]
    is a concatenation of right context, utterance and summary.

    Right context block is the concatenation of all the right context for
    each segments. [right_context_0, right_context_1, ..., right_context_n]
    For example, if we have utterance = [v0, v1, v2, ...., v20]. segment
    size 8, right_context size 4. Then the right context blocks =
    [v8, v9, v10, v11, v16, v17, v18, v19, 0, 0, 0, 0], where v8, v9, v10,
    and v11 are the right context for first segment. v16, v17, v18 and v19
    are the right context for second segment. 0, 0, 0 and 0 are right context
    for the last segment.

    utterance is corresponding to input embedding sequence

    summary is concatenation of average of each segments. [summary_0,
    summary_1, ..., ].

    In augmented memory multihead attention, the query is [right_context,
    utterance, summary], key is [memory, right_context, utterance]. Different
    with AugmentedMemoryMultiheadAttentionBmm, memory here is passed from
    previous attention layer. For the first attention layer, memory is average
    of each segment.

    Memory is a concatenation of memory from each segments in previous attention
    layer. For example, current layer is i, then memory is [m_0, m_1, ..., m_n].
    Each m_k is the output from seg_k in layer i-1.

    args:
        input_dim: input embedding dimension
        num_heads: number of heads in multihead self-attention
        dropout: attention dropout
        std_scale: if std_scale is not None. The weak attention suppression is
            turned on. For std_scale = 0.5, all the attention smaller than
            mean + 0.5 * std will be suppressed.
        scaled_init: whether to use scaled init for linear weight
        tanh_on_mem: whether to use tanh on memory output
        use_mem: whether to use memory or not. When max_memory_size is 0, then
            we don't have memory anymore.
        layer_index: current self-attention layer index that is used in depth
            initialization
        max_relative_position: max relative position used in relative position
            embedding
        rpe_old_option: To be compatible with previous model. The previous model
            was trained with attention += attention + rpe. The correct equation
            should be attention = attention + rpe

    """

    def __init__(
        self,
        input_dim,
        num_heads,
        dropout=0.0,
        std_scale=None,
        scaled_init=False,
        tanh_on_mem=False,
        use_mem=True,
        mini_batches=False,
        negative_inf="-inf",
        layer_index=-1,
        max_relative_position=0,
        rpe_old_option=True,
    ):
        if input_dim % num_heads:
            raise ValueError(
                "input_dim ({}) must be divisible by num_heads ({})".format(
                    input_dim, num_heads
                )
            )

        super().__init__()

        embed_dim = input_dim
        self.e2h_kv = torch.nn.Linear(input_dim, 2 * input_dim, bias=True)
        self.e2h_q = torch.nn.Linear(input_dim, input_dim, bias=True)
        self.rpe_old_option = rpe_old_option
        if max_relative_position > 0:
            self.use_rpe = True
            self.rpe_k = RelativePositionEmbedding(
                head_dim=input_dim // num_heads,
                max_position=max_relative_position,
            )
            self.rpe_v = RelativePositionEmbedding(
                head_dim=input_dim // num_heads,
                max_position=max_relative_position,
            )
        else:
            self.use_rpe = False
            self.rpe_k = None
            self.rpe_v = None
        if scaled_init:
            if layer_index == -1:
                gain = 1.0 / math.sqrt(2)
            else:
                # https://arxiv.org/abs/2005.09684 depthwise initialization
                # stablize the training greatly. Use depthwise initialization to
                # replace incremental loss.
                gain = 1.0 / math.sqrt(layer_index + 1)
            torch.nn.init.xavier_uniform_(self.e2h_kv.weight, gain=gain)
            torch.nn.init.xavier_uniform_(self.e2h_q.weight, gain=gain)

        self.out_proj = torch.nn.Linear(embed_dim, embed_dim, bias=True)

        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.dropout = dropout

        self.head_dim = embed_dim // num_heads
        self.scaling = self.head_dim ** -0.5

        self.std_scale = std_scale
        self.use_mem = use_mem
        self.mini_batches = mini_batches
        self.negative_inf = negative_inf

        if tanh_on_mem:
            self.squash_mem = torch.tanh
            self.nonlinear_squash_mem = True
        else:
            self.squash_mem = NoOp()
            self.nonlinear_squash_mem = False

    def prepare_qkv(
        self,
        input: Tensor,
        mems: Tensor,
        lengths: Tensor,
        summary_length: int,
        lc_length: int,
    ):
        # T: right_context length + utterance_length  + summary_length
        T, B, D = input.shape
        mem_length = mems.size(0)
        utterance_length = torch.max(lengths)

        right_context_blocks_length = T - utterance_length - summary_length
        rc_block = input[:right_context_blocks_length, :, :]
        utterance_block = input[right_context_blocks_length : T - summary_length, :, :]

        if B == 1:
            padding_mask = None
        else:
            klengths = lengths + mem_length + right_context_blocks_length + lc_length
            padding_mask = lengths_to_padding_mask(lengths=klengths)

        mem_rc_input = torch.cat([mems, rc_block, utterance_block], dim=0)

        # In training lc_length = 0
        key_length = mem_rc_input.size(0) + lc_length
        rc_input_sum = input
        q = self.e2h_q(rc_input_sum)
        kv = self.e2h_kv(mem_rc_input)
        k, v = kv.chunk(chunks=2, dim=2)
        result_qkv = (q, k, v)
        input_shape = (T, B, D)
        result_lengths_info = (
            mem_length,
            utterance_length,
            right_context_blocks_length,
            key_length,
        )
        if padding_mask is not None:
            assert padding_mask.size(0) == B
            assert padding_mask.size(1) == key_length

        return result_qkv, input_shape, result_lengths_info, padding_mask

    def prepare_attention_weights(
        self,
        q: Tensor,
        new_k: Tensor,
        new_v: Tensor,
        input_shape: Tuple[int, int, int],
        rpe: Optional[Tensor],
    ) -> Tuple[Tensor, Tensor, Tensor]:
        T, B, D = input_shape
        q = (
            q.contiguous().view(-1, B * self.num_heads, self.head_dim).transpose(0, 1)
            * self.scaling
        )

        k = (
            new_k.contiguous()
            .view(-1, B * self.num_heads, self.head_dim)
            .transpose(0, 1)
        )

        v = (
            new_v.contiguous()
            .view(-1, B * self.num_heads, self.head_dim)
            .transpose(0, 1)
        )

        attention_weights = torch.bmm(q, k.transpose(1, 2))
        if self.use_rpe and rpe is not None and self.rpe_v is not None:
            r_k = self.rpe_k(rpe)
            # [q, B*h, d] * [q, k, d] -> [B*h, q, k]
            attention_weights_rpe = torch.matmul(
                q.transpose(0, 1), r_k.transpose(1, 2)
            ).transpose(0, 1)
            attention_weights = attention_weights + attention_weights_rpe
        attention_weights_float = attention_weights.float()

        return attention_weights, attention_weights_float, v

    def prepare_attention_output(
        self,
        attention_weights: Tensor,
        attention_weights_float: Tensor,
        v: Tensor,
        input_shape: Tuple[int, int, int],
        key_length: int,
        padding_mask: Optional[Tensor],
        rpe: Optional[Tensor],
    ) -> Tensor:
        T, B, D = input_shape
        if padding_mask is not None:
            attention_weights_float = attention_weights_float.view(
                B, self.num_heads, T, key_length
            )
            attention_weights_float = attention_weights_float.masked_fill(
                padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), float("-inf")
            )
            attention_weights_float = attention_weights_float.view(
                B * self.num_heads, T, key_length
            )

        if self.std_scale is not None:
            attention_weights_float = attention_suppression(
                attention_weights_float, self.std_scale
            )

        attention_weights_float = torch.nn.functional.softmax(
            attention_weights_float, dim=-1
        )
        attention_weights = attention_weights_float.type_as(attention_weights)

        attention_probs = torch.nn.functional.dropout(
            attention_weights, p=self.dropout, training=self.training
        )

        # [T, key_length, B, n_head]+ [key_length, B, n_head, d_head]
        # -> [T, B, n_head, d_head]
        attention = torch.bmm(attention_probs, v)
        if self.use_rpe and rpe is not None and self.rpe_v is not None:
            r_v = self.rpe_v(rpe)
            attention_rpe = torch.matmul(
                attention_probs.transpose(0, 1), r_v
            ).transpose(0, 1)

            if self.rpe_old_option:
                attention += attention + attention_rpe
            else:
                attention = attention + attention_rpe

        assert list(attention.shape) == [B * self.num_heads, T, self.head_dim]

        attention = attention.transpose(0, 1).contiguous().view(T, B, self.embed_dim)

        rc_output_memory = self.out_proj(attention)
        return rc_output_memory

    @torch.jit.unused
    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

    @torch.jit.export
    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

    def quantize_(self, params=None):
        if params and "per_channel" in params and params["per_channel"]:
            qconfig = per_channel_dynamic_qconfig
        else:
            qconfig = default_dynamic_qconfig
        torch.quantization.quantize_dynamic(
            self, {torch.nn.Linear: qconfig}, dtype=torch.qint8, inplace=True
        )
        return self


class NoSegAugmentedMemoryTransformer(nn.Module):
    """
    Whole utterance augmented memory transformer.

    This is not pyspeech nn layer. It is used as a module in a master layer where
    multiple transformers is used.
    """

    def __init__(
        self,
        input_dim,
        num_heads,
        ffn_dim,
        dropout_in_attn=0.0,
        dropout_on_attn=None,
        dropout_on_fc1=None,
        dropout_on_fc2=None,
        activation_fn="relu",
        tanh_on_mem=False,
        std_scale=None,
        scaled_init=False,
        segment_size=128,
        use_mem=True,
        mini_batches=False,
        negative_inf="-inf",
        layer_index=-1,
        summarization_method="mean",
        max_relative_position=0,
        rpe_old_option=True,
    ):
        super(NoSegAugmentedMemoryTransformer, self).__init__()

        self.attention = NoSegAugmentedMemoryMultiheadAttentionBmm(
            input_dim=input_dim,
            num_heads=num_heads,
            dropout=dropout_in_attn,
            scaled_init=scaled_init,
            tanh_on_mem=tanh_on_mem,
            std_scale=std_scale,
            use_mem=use_mem,
            mini_batches=mini_batches,
            negative_inf=negative_inf,
            layer_index=layer_index,
            max_relative_position=max_relative_position,
        )
        self.dropout = nn.Dropout(dropout_on_attn)
        self.pos_ff = PositionwiseFF(
            input_dim=input_dim,
            ffn_dim=ffn_dim,
            dropout_on_fc1=dropout_on_fc1,
            dropout_on_fc2=dropout_on_fc2,
            activation_fn=activation_fn,
        )
        self.layer_norm_pre = Fp32LayerNorm(input_dim)
        self.layer_norm = Fp32LayerNorm(input_dim)
        self.segment_size = segment_size
        self.use_mem = use_mem

        self.memory_op = SummarizationLayer(
            summarization_method, segment_size, input_dim
        )

    def set_mini_batches(self, mini_batches):
        self.attention.mini_batches = mini_batches

    def gen_summary_queries(self, input):
        sum_input = self.memory_op(input)
        return sum_input

    def pre_attention_ops(self, input, right_context_blocks):
        rc_length = right_context_blocks.size(0)
        input_length = input.size(0)

        rc_and_input = torch.cat([right_context_blocks, input], dim=0)
        residual_input = rc_and_input
        rc_and_input = self.layer_norm_pre(rc_and_input)

        query_input = rc_and_input[-input_length:, :, :]
        return rc_length, input_length, residual_input, query_input, rc_and_input

    def after_attention_ops(self, attention_output, residual_input):
        output = self.dropout(attention_output)
        output = output + residual_input
        output = self.pos_ff(output)
        output = self.layer_norm(output)
        return output

    @torch.jit.export
    def forward_jit(
        self,
        input: Tensor,
        lengths: Tensor,
        mems: Tensor,
        left_context_key: Tensor,
        left_context_val: Tensor,
        right_context_blocks: Tensor,
        rpe: Optional[Tensor],
    ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]:

        results = self.pre_attention_ops(input, right_context_blocks)
        rc_length, input_length, residual_input, query_input, rc_and_input = results

        # In online decoding, the summary query size is always 1 or 0
        if self.use_mem:
            summary_query = self.gen_summary_queries(query_input)
            summary_query = summary_query[0:1, :, :]
            rc_qu_su = torch.cat([rc_and_input, summary_query], dim=0)
        else:
            rc_qu_su = rc_and_input

        rc_output, next_m, next_k, next_v = self.attention.forward_jit(
            input=rc_qu_su,
            lengths=lengths,
            mems=mems,
            left_context_key=left_context_key,
            left_context_val=left_context_val,
            rpe=rpe,
        )
        rc_output = self.after_attention_ops(rc_output, residual_input)
        results = (
            rc_output[-input_length:, :, :],
            next_m,
            rc_output[0:rc_length, :, :],
            next_k,
            next_v,
        )
        return results

    @torch.jit.unused
    def forward(
        self,
        input,
        lengths,
        mems,
        right_context_blocks,
        attention_mask,
        pre_mems,
        left_context_key,
        left_context_val,
        rpe,
    ):

        results = self.pre_attention_ops(input, right_context_blocks)
        rc_length, input_length, residual_input, query_input, rc_and_input = results
        if self.use_mem:
            summary_query = self.gen_summary_queries(query_input)
            rc_qu_su = torch.cat([rc_and_input, summary_query], dim=0)
        else:
            rc_qu_su = rc_and_input

        rc_output, next_m, next_k, next_v = self.attention(
            input=rc_qu_su,
            lengths=lengths,
            mems=mems,
            attention_mask=attention_mask,
            pre_mems=pre_mems,
            left_context_key=left_context_key,
            left_context_val=left_context_val,
            rpe=rpe,
        )

        # [TODO] Note memory did not go through pos_ff. What happen if we pass
        # memory through the pos_ff as well?
        rc_output = self.after_attention_ops(rc_output, residual_input)
        results = (
            rc_output[-input_length:, :, :],
            next_m,
            rc_output[0:rc_length, :, :],
            next_k,
            next_v,
        )

        return results


class NoSegAugmentedMemoryTransformerEncoderLayer(FairseqEncoder):
    """
    Whole utterance augmented memory transformer encoder layer. This is a master layer
    where we can define multiple augmented memory transformers. There are two reasons
    to setup the master layer.
    1. We only need to define once about the attention mask. All the layers in the master
       layer share the same mask.
    2. pyspeech nn layer has special input and output format. Defining one master layer is
       easier to passing memory between different layes inside the master layer

    args:
        input_dim: input embedding dimension
        num_heads: number of heads in multihead self-attention
        ffn_dim: ffn dimension in FFN layer
        num_layers: number of augmented memory transformer layers
        dropout_in_attn: dropout used in multi-head self-attention
        dropout_on_attn: dropout used for output from te multihead self-attention
        dropout_on_fc1: dropout used in FFN layer for the first linear layer
        dropout_on_fc2: dropout used in FFN layer for the second linear layer
        segment_size: segment size for each segment
        context_config: (left_context_size, right_context_size) defines the surround context size
            for each segment
        max_memory_size: maximum memory size used for each segment
        scaled_init: whether use scaled init for weight initialization in attention layer
        std_scale: if std_scale is not None. The weak attention suppression is
            turned on. For std_scale = 0.5, all the attention smaller than
            mean + 0.5 * std will be suppressed.
        activation_fn: activation function used in FFN layer. [ReLU, GELU] supported
        tanh_on_mem: whether use tanh on memory
        mini_batches: use mini-btach training
        negative_inf: the negative infinity value used in attention masking. default is "-inf".
            For some situation, e.g. LM. it is better to use "-1e8" to avoid nan issue.
        summarization_method: method to generate segment summrization embedding
        max_relative_position: max relatie position for relative position embedding
        rpe_old_option: To be compatible with previous model. The previous model
            was trained with attention += attention + rpe. The correct equation
            should be attention = attention + rpe
        [TODO]: remove the rpe_old_option by the end of 2021 Q1.

    """

    def __init__(
        self,
        input_dim,
        num_heads,
        ffn_dim,
        num_layers=1,
        dropout_in_attn=0.0,
        dropout_on_attn=0.0,
        dropout_on_fc1=0.0,
        dropout_on_fc2=0.0,
        segment_size=128,
        context_config=(0, 0),
        max_memory_size=0,
        scaled_init=True,
        std_scale=None,
        activation_fn="relu",
        tanh_on_mem=False,
        mini_batches=False,
        negative_inf="-inf",
        deep_init=True,
        summarization_method="mean",
        max_relative_position=0,
        rpe_old_option=True,
    ):
        super().__init__(None)
        if input_dim % num_heads:
            raise ValueError(
                "input_dim ({}) must be divisible by num_heads ({})".format(
                    input_dim, num_heads
                )
            )

        # we used to support growing memory size. However, it will cause
        # cross stream batching failure. Now we need to have exact max memory size
        if max_memory_size < 0:
            raise ValueError("max_memory_size must be >= 0")

        # Only assign right_context. In decoding, left context will be cached.
        # No need to let the online decoder to re-assign the left context
        self.left_context, self.right_context = context_config
        self.segment_size = segment_size
        self.memory_dim = input_dim
        self.max_memory_size = max_memory_size
        self.mini_batches = mini_batches
        if self.max_memory_size != 0:
            self.use_mem = True
        else:
            self.use_mem = False

        self.memory_op = SummarizationLayer(
            summarization_method, segment_size, input_dim
        )

        self.layers = torch.nn.ModuleList()
        self.num_layers = num_layers
        self.max_relative_position = max_relative_position
        if self.max_relative_position > 0:
            self.use_rpe = True
        else:
            self.use_rpe = False
        for i in range(self.num_layers):
            if deep_init:
                layer_index = i
            else:
                layer_index = -1

            self.layers.append(
                NoSegAugmentedMemoryTransformer(
                    num_heads=num_heads,
                    input_dim=input_dim,
                    ffn_dim=ffn_dim,
                    dropout_in_attn=dropout_in_attn,
                    dropout_on_attn=dropout_on_attn,
                    dropout_on_fc1=dropout_on_fc1,
                    dropout_on_fc2=dropout_on_fc2,
                    segment_size=segment_size,
                    std_scale=std_scale,
                    activation_fn=activation_fn,
                    tanh_on_mem=tanh_on_mem,
                    scaled_init=scaled_init,
                    use_mem=self.use_mem,
                    mini_batches=mini_batches,
                    negative_inf=negative_inf,
                    layer_index=layer_index,
                    summarization_method=summarization_method,
                    max_relative_position=max_relative_position,
                    rpe_old_option=rpe_old_option,
                )
            )

    def set_mini_batches(self, mini_batches):
        # handy function only used for unit test
        self.mini_batches = mini_batches
        for layer in self.layers:
            layer.set_mini_batches(mini_batches)

    def _get_relative_position(
        self,
        input: Tensor,
        max_relative_position: int,
        left_context_length: int,
        past_length: int,
        is_decoding: bool,
    ):
        # For training, we copy the right context to the start of the utterance
        # First dimension in distance is corresponding to query.
        # [right context, utterance, summary vector]
        # Second dimension in distance is corresponding to key.
        # [Memory bank, right context, utterance]
        # For summary vector in query part, the distance with
        # all other position is 2*max_position. For memory bank in key,
        # the distance with all other positions is 0.

        T, B, D = input.shape
        num_segs = math.ceil((T - self.right_context) / self.segment_size)

        # utterance
        u_st = past_length * self.segment_size
        u_ed = u_st + T
        utterance_ranges = torch.arange(u_st, u_ed - self.right_context)

        # left context. Only in minibatch or decoding
        left_context_ranges = torch.arange(u_st - left_context_length, u_st)

        # Right context block
        # right context + utterance
        right_context_blocks = []
        for i in range(0, num_segs - 1):
            st = (i + 1) * self.segment_size + u_st
            ed = st + self.right_context
            assert ed < u_ed
            temp = torch.arange(st, ed)
            right_context_blocks.append(temp)
        right_context_blocks.append(torch.arange(u_ed - self.right_context, u_ed))
        right_context_ranges = torch.cat(right_context_blocks)

        if self.use_mem:
            # Memory bank
            # The position for memory -n, .., -1
            if is_decoding:
                memory_size = min(past_length, self.max_memory_size)
            else:
                memory_size = num_segs + past_length - 1
            memory_bank_ranges = torch.arange(
                -max_relative_position - 1, -max_relative_position - 1 - memory_size, -1
            )

            # summary vector
            # The position for summary vector as the T+max_relative_position+1.
            # After the clamping, the relative position is max_relative_position
            summary_pos_st = u_ed + max_relative_position + 1
            summary_vector_ranges = torch.arange(
                summary_pos_st, summary_pos_st + num_segs
            )

            key_ranges = torch.cat(
                [
                    memory_bank_ranges,
                    right_context_ranges,
                    left_context_ranges,
                    utterance_ranges,
                ]
            )

            query_ranges = torch.cat(
                [right_context_ranges, utterance_ranges, summary_vector_ranges]
            )
        else:
            key_ranges = torch.cat(
                [right_context_ranges, left_context_ranges, utterance_ranges]
            )

            query_ranges = torch.cat([right_context_ranges, utterance_ranges])

        distance = key_ranges[None, :] - query_ranges[:, None]
        distance_clamp = (
            torch.clamp(distance, -max_relative_position, max_relative_position)
            + max_relative_position
        )
        distance_clamp = distance_clamp.to(input.device).long().detach()
        return distance_clamp

    def _get_attention_mask(self, input, past_length=0, left_context_cache=0):
        # attention mask for each query contains three parts:
        # 1. memory part
        # 2. left_context + segment
        # 3. right_context_block
        # so for each segment and its correspoinding right context block,
        # the attention matrix is formed by 9 parts:
        # [0, m, 0, 0, right_context, 0, 0, seg, 0]
        # [before memory, memory, after memory, before right context, right_context,
        #  after right context, before seg, seg, after seg]
        #
        # Query is formed in the way as [right_context_blocks, utterance, summary]
        #
        # Note: put m and right_context before segment is convenient
        # for padding_mask operation.
        # Key lengths = m_length + right_context_block_length + lengths
        utterance_length, batch_size, _ = input.shape
        summary_length = math.ceil(utterance_length / self.segment_size)
        num_segs = summary_length
        rc_length = self.right_context * num_segs
        rc = self.right_context
        lc = self.left_context

        # using mini-batches, there is left context cache available for current
        # sequence.
        lcc = left_context_cache

        # max_memory_size is 0 then we don't have memory and summary
        # past_length is the memory carry from previous sequence
        if self.use_mem:
            mem_length = num_segs - 1 + past_length
        else:
            mem_length = 0
        rc_mask = []
        query_mask = []
        summary_mask = []
        for j in range(0, num_segs):
            ssize = min(self.segment_size, utterance_length - j * self.segment_size)

            rc_size = rc
            rc_mat = []
            q_mat = []
            s_mat = []
            m_start = max(j + past_length - self.max_memory_size, 0)

            # max_memory_size is 0, then we don't use memory
            if self.use_mem:
                # part 0: before memory
                rc_mat.append(input.new_zeros(rc_size, m_start))
                q_mat.append(input.new_zeros(ssize, m_start))
                s_mat.append(input.new_zeros(1, m_start))

                # part 1: memory
                col_1 = j + past_length - m_start
                rc_mat.append(torch.ones(rc_size, col_1, device=input.device))
                q_mat.append(torch.ones(ssize, col_1, device=input.device))
                # based on D22875746, disable summary query attention
                # on memeory is better for long form utterance
                s_mat.append(input.new_zeros(1, col_1))

                # part 2: after memory
                col_2 = mem_length - (j + past_length)
                rc_mat.append(input.new_zeros(rc_size, col_2))
                q_mat.append(input.new_zeros(ssize, col_2))
                s_mat.append(input.new_zeros(1, col_2))

            # part 3: before right context
            rc_start = j * rc
            rc_mat.append(input.new_zeros(rc_size, rc_start))
            q_mat.append(input.new_zeros(ssize, rc_start))
            s_mat.append(input.new_zeros(1, rc_start))

            # part 4: right context
            rc_end = rc_start + rc
            col_4 = rc
            rc_mat.append(torch.ones(rc_size, col_4, device=input.device))
            q_mat.append(torch.ones(ssize, col_4, device=input.device))
            s_mat.append(torch.ones(1, col_4, device=input.device))

            # part 5: after right context
            col_5 = rc_length - rc_end
            rc_mat.append(input.new_zeros(rc_size, col_5))
            q_mat.append(input.new_zeros(ssize, col_5))
            s_mat.append(input.new_zeros(1, col_5))

            # part 6: before query segment
            seg_start = max(j * self.segment_size + lcc - lc, 0)
            rc_mat.append(input.new_zeros(rc_size, seg_start))
            q_mat.append(input.new_zeros(ssize, seg_start))
            s_mat.append(input.new_zeros(1, seg_start))

            # part 7: query segment
            # note: right context is put in right context block
            # here we only need to consider about left context
            seg_end = min((j + 1) * self.segment_size + lcc, utterance_length + lcc)
            col_7 = seg_end - seg_start
            rc_mat.append(torch.ones(rc_size, col_7, device=input.device))
            q_mat.append(torch.ones(ssize, col_7, device=input.device))
            s_mat.append(torch.ones(1, col_7, device=input.device))

            # part 8: after query segment
            col_8 = utterance_length + lcc - seg_end
            rc_mat.append(input.new_zeros(rc_size, col_8))
            q_mat.append(input.new_zeros(ssize, col_8))
            s_mat.append(input.new_zeros(1, col_8))

            rc_mask.append(torch.cat(rc_mat, dim=1))
            query_mask.append(torch.cat(q_mat, dim=1))
            summary_mask.append(torch.cat(s_mat, dim=1))

        # no memory, then we don't need summary either
        if self.use_mem:
            attention_mask = (
                1
                - torch.cat(
                    [
                        torch.cat(rc_mask, dim=0),
                        torch.cat(query_mask, dim=0),
                        torch.cat(summary_mask, dim=0),
                    ],
                    dim=0,
                )
            ).to(torch.bool)
        else:
            attention_mask = (
                1
                - torch.cat(
                    [torch.cat(rc_mask, dim=0), torch.cat(query_mask, dim=0)], dim=0
                )
            ).to(torch.bool)

        return attention_mask

    @torch.jit.export
    def init_state(
        self, batch_size: int, device: Optional[Device] = None
    ) -> List[Tensor]:
        empty_memory = torch.zeros(
            self.num_layers,
            self.max_memory_size,
            batch_size,
            self.memory_dim,
            device=device,
        )
        left_context_key = torch.zeros(
            self.num_layers,
            self.left_context,
            batch_size,
            self.memory_dim,
            device=device,
        )
        left_context_val = torch.zeros(
            self.num_layers,
            self.left_context,
            batch_size,
            self.memory_dim,
            device=device,
        )
        past_length = torch.zeros(1, batch_size, dtype=torch.int32, device=device)

        return [empty_memory, left_context_key, left_context_val, past_length]

    @torch.jit.export
    def batch_state(self, states: List[List[Tensor]]) -> List[Tensor]:
        if len(states) == 0:
            return []
        batched_m = []
        batched_lc_key = []
        batched_lc_val = []
        batched_past_length = []
        for state in states:
            if len(state) == 0:
                continue
            m, lc_key, lc_val, past_length = state
            batched_m.append(m)
            batched_lc_key.append(lc_key)
            batched_lc_val.append(lc_val)
            batched_past_length.append(past_length)

        if (
            (len(batched_m) == 0)
            or (len(batched_lc_key) == 0)
            or (len(batched_lc_val) == 0)
            or (len(batched_past_length) == 0)
        ):
            return [
                torch.tensor([]),
                torch.tensor([]),
                torch.tensor([]),
                torch.tensor([]),
            ]

        batched_m = torch.cat(batched_m, dim=2)
        batched_lc_key = torch.cat(batched_lc_key, dim=2)
        batched_lc_val = torch.cat(batched_lc_val, dim=2)
        batched_past_length = torch.cat(batched_past_length, dim=1)
        return [batched_m, batched_lc_key, batched_lc_val, batched_past_length]

    @torch.jit.export
    def reorder_state(self, state: List[Tensor], indices: Tensor) -> List[Tensor]:
        if len(state) == 0:
            return []
        m, lc_key, lc_val, past_length = state
        indices = indices.to(device=m.device)
        reord_m = torch.index_select(m, 2, indices)
        reord_lc_key = torch.index_select(lc_key, 2, indices)
        reord_lc_val = torch.index_select(lc_val, 2, indices)
        reord_past_length = torch.index_select(past_length, 1, indices)
        return [reord_m, reord_lc_key, reord_lc_val, reord_past_length]

    @torch.jit.export
    def reset_state(self, state: List[Tensor], indices: Tensor) -> List[Tensor]:
        m, lc_key, lc_val, past_length = state
        m = m.index_fill(dim=2, index=indices, value=0.0)
        lc_key = lc_key.index_fill(dim=2, index=indices, value=0.0)
        lc_val = lc_val.index_fill(dim=2, index=indices, value=0.0)
        past_length = past_length.index_fill(dim=1, index=indices, value=0)

        return [m, lc_key, lc_val, past_length]

    @torch.jit.export
    def state_size(self) -> int:
        return 4

    @torch.jit.export
    def batch_size_in_state(
        self, state: Optional[List[Tensor]], sloppy: bool = True
    ) -> Optional[int]:
        if state is None:
            return None
        return state[0].size(2)

    def gen_summary_queries(self, input):
        sum_input = self.memory_op(input)
        return sum_input

    def _gen_right_context_padded_input(self, input):
        # This function deals with input that is already
        # padded with right context (e.g. minibatch training)
        right_context_blocks = []
        T, B, D = input.shape
        num_segs = math.ceil((T - self.right_context) / self.segment_size)
        for i in range(0, num_segs - 1):
            st = (i + 1) * self.segment_size
            ed = st + self.right_context
            assert ed < T
            temp = input[st:ed, :, :]
            right_context_blocks.append(temp)

        # last segment right context is already available
        right_context_blocks.append(input[T - self.right_context :, :, :])
        return torch.cat(right_context_blocks, dim=0)

    def _gen_segs_right_context(self, input, lengths):
        segments = []
        T, B, D = input.size()
        nT = T - self.right_context

        # assume input is right context padded
        num_segs = math.ceil(nT / self.segment_size)
        # pad zeros to the utterance to make sure each
        # segment has the same right context. For the
        for i in range(0, num_segs - 1):
            st = i * self.segment_size
            ed = min(T, st + self.segment_size + self.right_context)
            temp = input[st:ed, :, :]
            rest_lengths = torch.clamp(
                lengths - self.segment_size, min=0, max=nT - (i + 1) * self.segment_size
            )
            segments.append((temp, lengths - rest_lengths + self.right_context))
            lengths = rest_lengths

        last_seg = input[st + self.segment_size :, :, :]
        segments.append((last_seg, rest_lengths + self.right_context))

        return segments

    @torch.jit.unused
    def forward(
        self, input: Tensor, padding_masks: Tensor, state: Optional[List[Tensor]] = None
    ) -> Tuple[Tensor, Tensor, List[Tensor], List[Tensor]]:
        # Xutai: originally the second argument is lengths.
        lengths = (~padding_masks).sum(dim=1).long()
        # mini batch training.
        if self.mini_batches:
            return self.forward_mini_batches(input, lengths, state)

        # regular full sequence training. Note, assume the right context in provided
        # in the input.
        T, B, D = input.size()
        right_context_blocks = self._gen_right_context_padded_input(input)

        # generate the relative positional embedding
        if self.use_rpe:
            rpe = self._get_relative_position(
                input=input,
                max_relative_position=self.max_relative_position,
                left_context_length=0,
                past_length=0,
                is_decoding=False,
            )
        else:
            rpe = None
        input = input[: T - self.right_context, :, :]

        attention_mask = self._get_attention_mask(input)

        # firt layer use each segment mean as memory
        # ignore the last one seg average
        if self.use_mem:
            mems = self.gen_summary_queries(input)[:-1, :, :]
        else:
            mems = torch.zeros(0, input.size(1), input.size(2), device=input.device)
            mems = mems.type_as(input)

        output = input
        all_outputs = []

        for layer in self.layers:
            output, mems, right_context_blocks, _, _ = layer(
                input=output,
                lengths=lengths,
                attention_mask=attention_mask,
                mems=mems,
                right_context_blocks=right_context_blocks,
                pre_mems=None,
                left_context_key=None,
                left_context_val=None,
                rpe=rpe,
            )
            all_outputs.append(output)
        return output, padding_masks, [], all_outputs

    def forward_jit_mini_batch_init(
        self,
        seg: Tensor,
        state: Optional[List[Tensor]] = None,
        is_decoding: bool = False,
    ):
        # Prepare state. In whole sequence training, state is ignored.
        # For minibatch training, we need to prepare state
        if state is None:
            state = self.init_state(batch_size=seg.size(1), device=seg.device)
            if seg.dtype == torch.half:
                state = [state[0].half(), state[1].half(), state[2].half(), state[3]]

        if self.use_mem:
            # note input average only on seg, not on right context
            # first layer use each segmetn mean as memory. the last
            # one segment average is used in state
            full_mems = self.gen_summary_queries(seg)
            if is_decoding:
                mems = full_mems[0:1, :, :]
                state_mems = torch.cat([state[0][0], mems], dim=0)
            else:
                mems = full_mems[:-1, :, :]
                state_mems = torch.cat([state[0][0], full_mems], dim=0)
        else:
            mems = state[0][0]
            state_mems = mems

        # track processed segment number or memory number
        # the same batch as the same bumber of past length
        past_length = state[3][0][0].item()
        past_left_context = min(past_length * self.segment_size, self.left_context)
        past_length = min(self.max_memory_size, past_length)

        return state, mems, state_mems, past_length, past_left_context

    def state_update_before(
        self, layer: int, state: List[Tensor], past_length: int, past_left_context: int
    ):
        pre_mems = state[0][layer][self.max_memory_size - past_length :, :, :]
        lc_key = state[1][layer][self.left_context - past_left_context :, :, :]
        lc_val = state[2][layer][self.left_context - past_left_context :, :, :]
        return pre_mems, lc_key, lc_val

    def state_update_after(
        self,
        layer: int,
        state: List[Tensor],
        mems: Tensor,
        next_key: Tensor,
        next_val: Tensor,
        mems_list: List[Tensor],
        lc_key_list: List[Tensor],
        lc_val_list: List[Tensor],
    ):
        # mems is used for next layer
        if layer < self.num_layers - 1:
            state_mems = torch.cat([state[0][layer + 1], mems], dim=0)
            mems_list.append(state_mems[-self.max_memory_size :, :, :])

        # when mems pass to next sequence, we need the last memory. when mems
        # use for the next layer, we can ignore the last memory
        mems = mems[:-1, :, :]

        # note state[1][i] and state[2][i] original length equals to self.left_context
        new_k = torch.cat([state[1][layer], next_key], dim=0)
        new_v = torch.cat([state[2][layer], next_val], dim=0)
        lc_key_list.append(new_k[-self.left_context :, :, :])
        lc_val_list.append(new_v[-self.left_context :, :, :])
        return mems_list, lc_key_list, lc_val_list, mems

    def state_update_after_loop(
        self,
        state: List[Tensor],
        mems_list: List[Tensor],
        lc_key_list: List[Tensor],
        lc_val_list: List[Tensor],
        update_length: int,
    ):
        state[0] = torch.stack(mems_list, dim=0)
        state[1] = torch.stack(lc_key_list, dim=0)
        state[2] = torch.stack(lc_val_list, dim=0)
        state[3] = state[3] + update_length
        return state

    @torch.jit.unused
    def forward_mini_batches(
        self, input: Tensor, lengths: Tensor, state: Optional[List[Tensor]] = None
    ) -> Tuple[Tensor, Tensor, List[Tensor], List[Tensor]]:
        T, B, D = input.size()

        # input without right context
        seg = input[: T - self.right_context, :, :]

        # get right context blocks
        right_context_blocks = self._gen_right_context_padded_input(input)

        mems_list = []
        lc_key_list = []
        lc_val_list = []
        results = self.forward_jit_mini_batch_init(seg, state, False)
        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=False,
            )
        else:
            rpe = None

        # get attention mask based on seg (not include right context) and available
        # left context
        attention_mask = self._get_attention_mask(seg, past_length, past_left_context)
        mems_list.append(state_mems[-self.max_memory_size :, :, :])
        output = seg
        i = 0
        all_outputs = []
        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.
            pre_mems, lc_key, lc_val = self.state_update_before(
                i, state, past_length, past_left_context
            )

            output, mems, right_context_blocks, next_key, next_val = layer.forward(
                input=output,
                lengths=lengths,
                attention_mask=attention_mask,
                mems=mems,
                right_context_blocks=right_context_blocks,
                pre_mems=pre_mems,
                left_context_key=lc_key,
                left_context_val=lc_val,
                rpe=rpe,
            )
            all_outputs.append(output)
            mems_list, lc_key_list, lc_val_list, mems = self.state_update_after(
                layer=i,
                state=state,
                mems=mems,
                next_key=next_key,
                next_val=next_val,
                mems_list=mems_list,
                lc_key_list=lc_key_list,
                lc_val_list=lc_val_list,
            )

            i += 1

        # update state
        update_length = math.ceil((T - self.right_context) / self.segment_size)
        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=update_length,
        )

        return output, lengths, state, all_outputs

    def forward_jit_test(
        self, input: Tensor, lengths: Tensor, state: Optional[List[Tensor]] = None
    ) -> Tuple[Tensor, Tensor, List[Tensor]]:
        """
        This one simulate sequence encoder forward jit. This is for unit test purpose.
        It is not used in training or decoding. Note, extra_right_context is set in
        the model. In unit test, input = [utterance, right_context], lengths =
        [utterance_length].
        args:
            input: input utterance
            lengths: utterance input length
            state: None here. input is whole utterance
        """
        # [TODO] sequence_to_segment has bug in lengths.
        seg_src_tokens_lengths = self._gen_segs_right_context(input, lengths)

        seg_enc_tokens_lengths: List[Tuple[Tensor, Tensor]] = []
        state: Optional[List[Tensor]] = None
        for seg_src_tokens, seg_src_lengths in seg_src_tokens_lengths:
            seg_enc_tokens, seg_enc_lengths, state = self.forward_jit(
                input=seg_src_tokens, lengths=seg_src_lengths, state=state
            )
            seg_enc_tokens_lengths.append((seg_enc_tokens, seg_enc_lengths))

        enc_tokens, enc_lengths = segments_to_sequence(
            segments=seg_enc_tokens_lengths, time_axis=0
        )

        state = []  # returns trivial state

        return enc_tokens, enc_lengths, state

    @torch.jit.export
    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

    def quantize_(self, params=None):
        if params and "per_channel" in params and params["per_channel"]:
            qconfig = per_channel_dynamic_qconfig
        else:
            qconfig = default_dynamic_qconfig
        torch.quantization.quantize_dynamic(
            self, {torch.nn.Linear: qconfig}, dtype=torch.qint8, inplace=True
        )
        return self


# ------------------------------------------------------------------------------
#   Emformer encoder for seq2seq model
#   This is a wrapper over the original emformer
# ------------------------------------------------------------------------------
def emformer_encoder(klass):
    class SpeechEncoder(klass):
        def __init__(self, args):
            super().__init__(args)
            stride = SpeechEncoder.conv_layer_stride(args)
            trf_left_context = args.segment_left_context // stride
            trf_right_context = args.segment_right_context // stride
            context_config = [trf_left_context, trf_right_context]
            self.transformer_layers = nn.ModuleList(
                [
                    NoSegAugmentedMemoryTransformerEncoderLayer(
                        input_dim=args.encoder_embed_dim,
                        num_heads=args.encoder_attention_heads,
                        ffn_dim=args.encoder_ffn_embed_dim,
                        num_layers=args.encoder_layers,
                        dropout_in_attn=args.dropout,
                        dropout_on_attn=args.dropout,
                        dropout_on_fc1=args.dropout,
                        dropout_on_fc2=args.dropout,
                        activation_fn=args.activation_fn,
                        context_config=context_config,
                        segment_size=args.segment_length,
                        max_memory_size=args.max_memory_size,
                        scaled_init=True,  # TODO: use constant for now.
                        tanh_on_mem=args.amtrf_tanh_on_mem,
                    )
                ]
            )

        def forward(self, src_tokens, src_lengths):
            encoder_out = super().forward(src_tokens, src_lengths)
            output = encoder_out["encoder_out"][0]
            encoder_padding_masks = encoder_out["encoder_padding_mask"][0]

            # This is because that in the original implementation
            # the output didn't consider the last segment as right context.
            encoder_padding_masks = encoder_padding_masks[:, : output.size(0)]

            return {
                "encoder_out": [output],
                "encoder_padding_mask": [encoder_padding_masks],
                "encoder_embedding": [],
                "encoder_states": [],
                "src_tokens": [],
                "src_lengths": [],
            }

        @staticmethod
        def conv_layer_stride(args):
            # TODO: make it configurable from the args
            return 4

    SpeechEncoder.__name__ = klass.__name__
    return SpeechEncoder