class_docstr: Transformer Encoder Layer
docstr: Parameters
  ----------
  data :
      If layout == 'NT'
          Shape (batch_size, seq_length, C_in)
      Else
          Shape (seq_length, batch_size, C_in)
  attn_mask :
      Shape (batch_size, seq_length, seq_length)
  Returns
  -------
  out :
      If layout == 'NT'
          Shape (batch_size, seq_length, C_out)
      Else
          Shape (seq_length, batch_size, C_out)
  attn_weight :
      Shape (batch_size, seq_length, seq_length)
docstr_tok: ["Parameters","--","--","--","--","--","data",":","If","layout","==","'NT'","Shape","(","batch_size",",","seq_length",",","C_in",")","Else","Shape","(","seq_length",",","batch_size",",","C_in",")","attn_mask",":","Shape","(","batch_size",",","seq_length",",","seq_length",")","Returns","--","--","--","-","out",":","If","layout","==","'NT'","Shape","(","batch_size",",","seq_length",",","C_out",")","Else","Shape","(","seq_length",",","batch_size",",","C_out",")","attn_weight",":","Shape","(","batch_size",",","seq_length",",","seq_length",")"]
code: def forward(self, data, attn_mask):
      """
      Parameters
      ----------
      data :
          If layout == 'NT'
              Shape (batch_size, seq_length, C_in)
          Else
              Shape (seq_length, batch_size, C_in)
      attn_mask :
          Shape (batch_size, seq_length, seq_length)
      Returns
      -------
      out :
          If layout == 'NT'
              Shape (batch_size, seq_length, C_out)
          Else
              Shape (seq_length, batch_size, C_out)
      attn_weight :
          Shape (batch_size, seq_length, seq_length)
      """
      residual = data
      if self._pre_norm:
          data = self.layer_norm(data)
      query, key, value = th.split(
          self.attn_qkv(data), self._units, dim=-1
      )
      query = th.reshape(query, query.shape[:2] + (self._num_heads, -1))
      key = th.reshape(key, key.shape[:2] + (self._num_heads, -1))
      value = th.reshape(value, value.shape[:2] + (self._num_heads, -1))
      out, [_, attn_weight] = self.attention_cell(
          query, key, value, attn_mask
      )
      out = self.attention_proj(out)
      out = self.dropout_layer(out)
      out = out + residual
      if not self._pre_norm:
          out = self.layer_norm(out)
      out = self.ffn(out)
      return out, attn_weight
  
imports: from typing import Union, Optional, List
  from torch import nn
  import torch
  from ..layers import PositionwiseFFN
  from ..attention_cell import MultiHeadAttentionCell, gen_mem_attn_mask, gen_self_attn_mask
  from ..utils import to_torch_dtype
label_desc: Use image node_id 3 for calling the TransformerEncoderLayer obj's underlying member method code with example usage: obj.forward(data, attn_mask) and returns: out,  attn_weight