code: def ReformerLM(
      vocab_size,
      d_model=512,
      d_ff=2048,
      d_attention_key=64,
      d_attention_value=64,
      n_layers=6,
      n_heads=8,
      dropout=0.1,
      max_len=2048,
      attention_type=tl.SelfAttention,
      pos_type=None,
      pos_axial_shape=(),
      pos_d_axial_embs=None,
      pos_start_from_zero_prob=1.0,
      pos_max_offset_to_add=0,
      ff_activation=tl.FastGelu,
      ff_use_sru=0,
      ff_chunk_size=0,
      ff_sparsity=0,
      loss_sparsity_type="mult",
      loss_sparsity=0,
      loss_d_lowrank=0,
      loss_sparsity_prob=None,
      attention_chunk_size=0,
      mode="train",
  ):
      """Reversible transformer language model (only uses a decoder, no encoder).
  
      Args:
        vocab_size: int: vocab size
        d_model: int:  depth of *each half* of the two-part features
        d_ff: int: depth of feed-forward layer
        d_attention_key: int: depth of key vector for each attention head
        d_attention_value: int: depth of value vector for each attention head
        n_layers: int: number of decoder layers
        n_heads: int: number of attention heads
        dropout: float: dropout rate (how much to drop out)
        max_len: int: maximum symbol length for positional encoding
        attention_type: class: attention class to use, such as SelfAttention.
        pos_type: string, the type of positional embeddings to use.
        pos_axial_shape: tuple of ints: input shape to use for the axial position
          encoding. If unset, axial position encoding is disabled.
        pos_d_axial_embs: tuple of ints: depth of position embedding for each axis.
          Tuple length must match pos_axial_shape, and values must sum to d_model.
        pos_start_from_zero_prob: how often to start from 0 during training,
              (if 1.0, we always start from position 0, if less, we randomize).
        pos_max_offset_to_add: maximum offset to add to positions during training
            when randomizing; this offset plus input length must still be less than
            max_len for all training examples.
        ff_activation: the non-linearity in feed-forward layer
        ff_use_sru: int; if > 0, we use this many SRU layers instead of feed-forward
        ff_chunk_size: int; if > 0, chunk feed-forward into this-sized chunks
        ff_sparsity: int, if > 0 use sparse feed-forward block with this sparsity
        loss_sparsity_type: str, type of sparsity to used in loss layer. See
          SparseDenseWithOptions for options. None if no sparsity should be used.
        loss_sparsity: int, the sparsity for loss layer (if used)
        loss_d_lowrank: int, the dimensions for intermediate layer (if used)
        loss_sparsity_prob: float, the probability for sparse version of loss to be
          used. If None, only sparse version is used.
        attention_chunk_size: int, if > 0 run attention chunked at this size
        mode: str: 'train', 'eval', or 'predict'
  
      Returns:
        the layer.
      """
      positional_encoding = ct.PositionalEncoder(
          mode,
          dropout,
          max_len,
          pos_type,
          pos_axial_shape,
          pos_d_axial_embs,
          pos_start_from_zero_prob,
          pos_max_offset_to_add,
      )
  
      positional_embedder = [
          tl.Embedding(vocab_size, d_model),
          tl.Dropout(
              rate=dropout, shared_axes=[-2], mode=mode
          ),  # pylint: disable=no-value-for-parameter
          positional_encoding,
      ]
  
      decoder_blocks = []
  
      if isinstance(attention_type, (tuple, list)):
          assert n_layers % len(attention_type) == 0
      else:
          attention_type = [attention_type]
      for layer_idx in range(n_layers):
          layer_attention_type = attention_type[
              layer_idx % len(attention_type)
          ]
          decoder_block = DecoderBlock(
              d_model,
              d_ff,
              d_attention_key,
              d_attention_value,
              n_heads,
              attention_type=layer_attention_type,
              dropout=dropout,
              ff_activation=ff_activation,
              ff_dropout=dropout,
              ff_use_sru=ff_use_sru,
              ff_chunk_size=ff_chunk_size,
              ff_sparsity=ff_sparsity,
              attention_chunk_size=attention_chunk_size,
              mode=mode,
          )
          decoder_blocks.append(decoder_block)
  
      dense_loss_layer = tl.SparseDenseWithOptions(
          vocab_size,
          d_input=d_model,
          sparsity_type=loss_sparsity_type,
          sparsity=loss_sparsity,
          d_lowrank=loss_d_lowrank,
          prob_sparse=loss_sparsity_prob,
          mode=mode,
      )
  
      return tl.Serial(
          tl.ShiftRight(mode=mode),
          positional_embedder,
          tl.Dup(),
          tl.ReversibleSerial(decoder_blocks),
          tl.Concatenate(),
          # TODO(kitaev): Test whether dropout should go before or after the
          # LayerNorm, and whether dropout broadcasting is needed here.
          tl.LayerNorm(),
          tl.Dropout(
              rate=dropout, shared_axes=[-2], mode=mode
          ),  # pylint: disable=no-value-for-parameter
          dense_loss_layer,
      )
  
file_path: trax/trax/models/reformer/reformer.py
imports: from trax import layers
  from trax.fastmath import numpy
  from trax.models.research import configurable_transformer
label_desc: Use image node_id 2 for calling a global function with example usage: ReformerLM(vocab_size, d_model, d_ff, d_attention_key, d_attention_value, n_layers, n_heads, dropout, max_len, attention_type, pos_type, pos_axial_shape, pos_d_axial_embs, pos_start_from_zero_prob, pos_max_offset_to_add, ff_activation, ff_use_sru, ff_chunk_size, ff_sparsity, loss_sparsity_type, loss_sparsity, loss_d_lowrank, loss_sparsity_prob, attention_chunk_size, mode) and returns: tl