def forward(
            self,
            inputs: torch.Tensor,
            pos_embedding: torch.Tensor,
            u: torch.nn.Parameter,
            v: torch.nn.Parameter,
            memory: torch.Tensor,
            mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        **Overview**:
            The forward computation graph of GTrXL layer.
        """
        # Concat memory with input across sequence dimension. The shape is: [full_sequence, batch_size, input_dim]
        full_input = torch.cat([memory, inputs], dim=0)
        # Forward calculation for GTrXL layer.
        # In GTrXL, the layer normalization is put before the attention layer.
        x1 = self.layernorm1(full_input)
        # Attention module.
        a1 = self.dropout(self.attention(inputs, pos_embedding, x1, u, v, mask=mask))
        a1 = self.activation(a1)
        # In GTrXL, gating layer replace the resnet layer in TrXL.
        o1 = self.gate1(inputs, a1) if self.gating else inputs + a1


        x2 = self.layernorm2(o1)
        # Feed Forward Network.
        m2 = self.dropout(self.mlp(x2))
        o2 = self.gate2(o1, m2) if self.gating else o1 + m2
        return o2

        **Overview**:
            The forward computation graph of GTrXL layer.
        
chapter5_time/gtrxl.py