def forward(self, x: torch.Tensor, batch_first: bool = False, return_mem: bool = True) -> Dict[str, torch.Tensor]:
        """
        **Overview**:
            The forward computation graph of GTrXL.
        """
        # If the first dimension of input x is batch_size,
        # then reshape x from  [batch_size ,sequence_length ,input_dim] to [sequence_length, batch_size, input_dim]
        if batch_first:
            x = torch.transpose(x, 1, 0)
        cur_seq, bs = x.shape[:2]
        # Get back memory.
        memory = None if self.memory is None else self.memory.get()
        # Abnormal case: no memory or memory shape mismatch.
        if memory is None:
            self.reset_memory(bs)
        elif memory.shape[-2] != bs or memory.shape[-1] != self.embedding_dim:
            warnings.warn(
                "Memory {} and Input {} dimensions don't match,"
                " this will cause the memory to be initialized to fit your input!".format(
                    list(memory.shape[-2:]), [x.shape[-2]] + [self.embedding_dim]
                )
            )
            self.reset_memory(bs)
        self.memory.to(x.device)
        memory = self.memory.get()
        # Pass through embedding layer.
        if self.use_embedding_layer:
            x = self.dropout(self.embedding(x))
        # Get full sequence length: memory length + current length
        prev_seq = self.memory_len
        full_seq = cur_seq + prev_seq
        # If the attention mask for current sequence length is already created, reuse the mask stored in ``self.att_mask`` .
        if cur_seq in self.att_mask.keys():
            attn_mask = self.att_mask[cur_seq]
        # Otherwise, create a new attention mask and store it into ``self.att_mask`` .
        else:
            # For example, if cur_seq = 3, full_seq = 7, then the mask is:
            # $$ \begin{matrix} 0 & 0 & 0 & 0 & 0 & 1 & 1 \\ 0 & 0 & 0 & 0 & 0 & 0 & 1 \\ 0 & 0 & 0 & 0 & 0 & 0 & 0 \end{matrix}$$
            # This forces that the hidden state of current token is only associated with previous tokens.
            attn_mask = (
                torch.triu(
                    torch.ones((cur_seq, full_seq)),
                    diagonal=1 + prev_seq,
                ).bool().unsqueeze(-1).to(x.device)
            )
            self.att_mask[cur_seq] = attn_mask
        # If the position encoding for current sequence length is already created, reuse it stored in ``self.pos_embedding_dict`` .
        if cur_seq in self.pos_embedding_dict.keys():
            pos_embedding = self.pos_embedding_dict[cur_seq]
        # Otherwise, create a new position encoding and store it into ``self.pos_embedding_dict`` .
        else:
            pos_ips = torch.arange(full_seq - 1, -1, -1.0, dtype=torch.float)  # full_seq
            pos_embedding = self.pos_embedding(pos_ips.to(x.device))
            self.pos_embedding_dict[cur_seq] = pos_embedding
        pos_embedding = self.dropout(pos_embedding)  # full_seq x 1 x embedding_dim

        hidden_state = [x]
        out = x
        # Calculate results for each GTrXL layer.
        for i in range(self.layer_num):
            layer = self.layers[i]
            out = layer(
                out,
                pos_embedding,
                self.u,
                self.v,
                mask=attn_mask,
                memory=memory[i],
            )
            hidden_state.append(out.clone())
        out = self.dropout(out)
        # Update the GTrXL memory.
        self.memory.update(hidden_state)
        # If the first dimension of output is required to be batch_size, then reshape x from  [sequence_length, batch_size, input_dim] to [batch_size ,sequence_length ,input_dim].
        if batch_first:
            out = torch.transpose(out, 1, 0)
        # Return memory is needed.
        if return_mem:
            output = treetensor.Object({"logit": out, "memory": memory})
        else:
            output = treetensor.Object({"logit": out})
        return output

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