code: def RNNLM(
      vocab_size,
      d_model=512,
      n_layers=2,
      rnn_cell=tl.LSTMCell,
      rnn_cell_d_state_multiplier=2,
      dropout=0.1,
      mode="train",
  ):
      """Returns an RNN language model.
  
      This model performs autoregressive language modeling:
  
        - input: rank 2 tensor representing a batch of text strings via token IDs
          plus padding markers; shape is (batch_size, sequence_length). The tensor
          elements are integers in `range(vocab_size)`, and `0` values mark padding
          positions.
  
        - output: rank 3 tensor representing a batch of log-probability
          distributions for each sequence position over possible token IDs;
          shape is (batch_size, sequence_length, `vocab_size`).
  
      Args:
        vocab_size: Input vocabulary size -- each element of the input tensor
            should be an integer in `range(vocab_size)`. These integers typically
            represent token IDs from a vocabulary-based tokenizer.
        d_model: Embedding depth throughout the model.
        n_layers: Number of RNN layers.
        rnn_cell: Type of RNN cell; must be a subclass of `Layer`.
        rnn_cell_d_state_multiplier: Multiplier for feature depth of RNN cell
            state.
        dropout: Stochastic rate (probability) for dropping an activation value
            when applying dropout.
        mode: If `'predict'`, use fast inference; if `'train'` apply dropout.
  
      Returns:
        An RNN language model as a layer that maps from a tensor of tokens
        to activations over a vocab set.
      """
  
      if (
          n_layers != 2
      ):  # TODO(jonni): Remove n_layers arg, if it can't vary?
          raise ValueError(
              f"Number of layers must be set to 2; instead got"
              f" {n_layers}."
          )
  
      def MultiRNNCell():
          """Multi-layer RNN cell."""
          return tl.Serial(
              tl.Parallel([], tl.Split(n_items=n_layers)),
              tl.SerialWithSideOutputs(
                  [rnn_cell(n_units=d_model) for _ in range(n_layers)]
              ),
              tl.Parallel([], tl.Concatenate(n_items=n_layers)),
          )
  
      zero_state = (
          tl.MakeZeroState(  # pylint: disable=no-value-for-parameter
              depth_multiplier=n_layers * rnn_cell_d_state_multiplier
          )
      )
  
      return tl.Serial(
          tl.ShiftRight(mode=mode),
          tl.Embedding(vocab_size, d_model),
          tl.Dropout(rate=dropout, mode=mode),
          tl.Branch([], zero_state),
          tl.Scan(MultiRNNCell(), axis=1, mode=mode),
          tl.Select([0], n_in=2),  # Drop RNN state.
          tl.Dense(vocab_size),
      )
  
imports: from trax import layers
  from trax.fastmath import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: RNNLM(vocab_size, d_model, n_layers, rnn_cell, rnn_cell_d_state_multiplier, dropout, mode) and returns: tl, tl