docstr: Generate the representation given the inputs.
  
  This is used in training or fine-tuning a mobile bert model.
  
  Parameters
  ----------
  data
      - layout = 'NT'
          Shape (batch_size, seq_length, C)
      - layout = 'TN'
          Shape (seq_length, batch_size, C)
  
  valid_length
      Shape (batch_size,)
  
  Returns
  -------
  out
      - layout = 'NT'
          Shape (batch_size, seq_length, C_out)
      - layout = 'TN'
          Shape (seq_length, batch_size, C_out)
docstr_tok: ["Generate","the","representation","given","the","inputs",".","This","is","used","in","training","or","fine-tuning","a","mobile","bert","model",".","Parameters","--","--","--","--","--","data","-","layout","=","'NT'","Shape","(","batch_size",",","seq_length",",","C",")","-","layout","=","'TN'","Shape","(","seq_length",",","batch_size",",","C",")","valid_length","Shape","(","batch_size",",",")","Returns","--","--","--","-","out","-","layout","=","'NT'","Shape","(","batch_size",",","seq_length",",","C_out",")","-","layout","=","'TN'","Shape","(","seq_length",",","batch_size",",","C_out",")"]
code: def forward(self, data, valid_length):
      """
      Generate the representation given the inputs.
  
      This is used in training or fine-tuning a mobile bert model.
  
      Parameters
      ----------
      data
          - layout = 'NT'
              Shape (batch_size, seq_length, C)
          - layout = 'TN'
              Shape (seq_length, batch_size, C)
  
      valid_length
          Shape (batch_size,)
  
      Returns
      -------
      out
          - layout = 'NT'
              Shape (batch_size, seq_length, C_out)
          - layout = 'TN'
              Shape (seq_length, batch_size, C_out)
      """
      if self._layout == "NT":
          batch_axis, time_axis = 0, 1
      elif self._layout == "TN":
          batch_axis, time_axis = 1, 0
      else:
          raise NotImplementedError(
              'Received layout="{}". '
              'Only "NT" and "TN" are supported.'.format(self._layout)
          )
      # 1. Embed the data
      attn_mask = gen_self_attn_mask(
          data,
          valid_length,
          dtype=self._dtype,
          layout=self._layout,
          attn_type="full",
      )
      out = data
      all_encodings_outputs = []
      additional_outputs = []
      all_encodings_outputs.append(out)
      for layer_idx in range(self._num_layers):
          layer = self.all_layers[layer_idx]
          out, attention_weights = layer(out, attn_mask)
          # out : [batch_size, seq_len, units]
          # attention_weights : [batch_size, num_heads, seq_len, seq_len]
          if self._output_all_encodings:
              out = npx.sequence_mask(
                  out,
                  sequence_length=valid_length,
                  use_sequence_length=True,
                  axis=time_axis,
              )
              all_encodings_outputs.append(out)
  
          if self._output_attention:
              additional_outputs.append(attention_weights)
  
      if not self._output_all_encodings:
          # if self._output_all_encodings, SequenceMask is already applied above
          out = npx.sequence_mask(
              out,
              sequence_length=valid_length,
              use_sequence_length=True,
              axis=time_axis,
          )
          return out, additional_outputs
      else:
          return all_encodings_outputs, additional_outputs
  
imports: import os
  from typing import Tuple, Optional
  import mxnet
  from mxnet import use_np, np, npx
  from mxnet.gluon import HybridBlock, nn
  from .base import BACKBONE_REGISTRY
  from ..op import select_vectors_by_position
  from ..base import get_model_zoo_home_dir, get_repo_model_zoo_url, get_model_zoo_checksum_dir
  from ..layers import InitializerType, PositionwiseFFN, PositionalEmbedding, get_norm_layer, get_activation
  from ..initializer import TruncNorm
  from ..utils.config import CfgNode
  from ..utils.misc import load_checksum_stats, download
  from ..utils.registry import Registry
  from ..attention_cell import MultiHeadAttentionCell, gen_self_attn_mask
  from ..data.tokenizers import HuggingFaceWordPieceTokenizer
label_desc: Use image node_id 3 for calling the MobileBertTransformer obj's underlying member method code with example usage: obj.forward(data, valid_length) and returns: out,  additional_outputs, all_encodings_outputs,  additional_outputs