code: def ResidualSwitchUnit(
      d_model, dropout=0.1, mode="train", residual_weight=0.9
  ):
      r"""RSU (Residual Switch Unit) layer as in https://arxiv.org/pdf/2004.04662.pdf.
  
    As defined in the paper:
  
    .. math::
      i &= [i_1, i_2] \\
      g &= GELU(LayerNorm(Z i)) \\
      c &= W g + B \\
      [o_1, o_2] &= \sigma(S) \bigodot i + h \bigodot c
  
    where Z, W, B, S are learnable parameters with sizes 2m × 4m, 4m × 2m, 2m, 2m.
    We assume that both i_1 and i_2 have size m. h is a scalar value.
  
    We assume the input is of shape [batch, length, depth].
  
    Args:
      d_model: output depth of the SRU layer
      dropout: dropout rate used in 'train' mode
      mode: mode for dropout layer
      residual_weight: value used in initializing vector S and constant h
  
    Returns:
      The RSU layer.
    """
      return tl.Serial(
          tl.Fn(
              "Reshape2Pairs",
              lambda x: jnp.reshape(
                  x, (x.shape[0], x.shape[1] // 2, -1)
              ),
              n_out=1,
          ),
          tl.Residual(
              tl.Dense(4 * d_model, use_bias=False),
              tl.LayerNorm(),
              tl.Gelu(),
              tl.Dense(2 * d_model),
              tl.Fn(
                  "Scaling",
                  lambda x: x
                  * np.sqrt(1 - residual_weight**2)
                  * 0.25,
                  n_out=1,
              ),
              shortcut=_ClippedScaling(residual_weight),
          ),
          tl.Fn(
              "UnPair",
              lambda x: jnp.reshape(
                  x, (x.shape[0], x.shape[1] * 2, -1)
              ),
              n_out=1,
          ),
          tl.Dropout(rate=dropout, mode=mode),
      )
  
imports: import numpy
  from trax import fastmath
  from trax import layers
  from trax.fastmath import numpy
  from trax.layers.assert_shape import assert_shape
label_desc: Use image node_id 2 for calling a global function with example usage: ResidualSwitchUnit(d_model, dropout, mode, residual_weight) and returns: tl