code: def SkippingTransformerLM(
      vocab_size,
      d_model=512,
      d_ff=2048,
      n_layers=6,
      n_heads=8,
      dropout=0.1,
      max_len=2048,
      mode="train",
      ff_activation=tl.Relu,
      skip_fraction=0.4,
  ):
      """Returns a Skipping Transformer language model.
  
      The input to the model is a tensor of tokens. (This model uses only the
      decoder part of the overall Transformer.)
  
      Args:
        vocab_size: int: vocab size
        d_model: int:  depth of embedding
        d_ff: int: depth of feed-forward layer
        n_layers: int: number of encoder/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
        mode: str: 'train', 'eval' or 'predict', predict mode is for fast inference
        ff_activation: the non-linearity in feed-forward layer
        skip_fraction: fraction of times to skip some layers
  
      Returns:
        A Transformer language model as a layer that maps from a tensor of tokens
        to activations over a vocab set.
      """
      embedder = [
          tl.Embedding(vocab_size, d_model),
          tl.Dropout(rate=dropout, mode=mode),
          tl.PositionalEncoding(max_len=max_len, mode=mode),
      ]
  
      @assert_shape("...sd,->...sd,")
      def ConditionedBlock(current_layer_num):
          return tl.Serial(
              # stack: embedding, n_layers_to_keep
              tl.Select(
                  [1, 0, 1]
              ),  # n_layers_to_keep, embedding, n_layers_to_keep
              tl.Cond(
                  # if n_layers_to_keep > current_layer_num
                  LargerThan(float(current_layer_num)),
                  # then: run block
                  tl.Serial(
                      transformer._DecoderBlock(  # pylint: disable=g-complex-comprehension,protected-access
                          d_model,
                          d_ff,
                          n_heads,
                          dropout,
                          [],
                          mode,
                          ff_activation,
                      )
                  ),
                  # else: run noop
                  tl.Serial(),
              )
              # stack: embedding, n_layers_to_keep
          )
  
      if mode == "train":
          if skip_fraction == 0.0:
              minimum_layers = float(n_layers)
              maximum_layers = float(n_layers)
          else:
              minimum_layers = 0.0
              maximum_layers = float(n_layers) / skip_fraction
      else:
          minimum_layers = maximum_layers = float(n_layers)
  
      return tl.Serial(
          tl.ShiftRight(mode=mode),
          embedder,
          # stack: embedding
          tl.RandomUniform(minimum_layers, maximum_layers, sync=True),
          # stack: n_layers_to_keep, embedding
          tl.Swap(),
          # stack: embedding, n_layers_to_keep
          [ConditionedBlock(i) for i in range(n_layers)],
          # stack: embedding, n_layers_to_keep
          tl.AssertShape("...sd,"),
          tl.Select([0], n_in=2),  # stack: embedding
          tl.AssertShape("...sd"),
          tl.LayerNorm(),
          tl.Dense(vocab_size),
      )
  
imports: from trax import layers
  from trax.layers.assert_shape import assert_shape
  from trax.models import transformer
label_desc: Use image node_id 2 for calling a global function with example usage: SkippingTransformerLM(vocab_size, d_model, d_ff, n_layers, n_heads, dropout, max_len, mode, ff_activation, skip_fraction) and returns: tl, tl