def __call__(self, inputs, deterministic, attn_mask=None):
    """Applies Transformer module.

    Args:
      inputs: input data.
      deterministic: deterministic
      attn_mask: self-attention mask.

    Returns:
      output from VTNLayer: [batch_size, seq_length, model_dim].
    """
    config = self.config

    # Attention block.
    assert inputs.ndim == 3
    x = nn.LayerNorm(dtype=config.dtype)(inputs)
    x = nn.SelfAttention(
        num_heads=config.num_heads,
        dtype=config.dtype,
        qkv_features=config.qkv_dim,
        kernel_init=config.kernel_init,
        bias_init=config.bias_init,
        use_bias=False,
        broadcast_dropout=False,
        dropout_rate=config.attention_dropout_rate,
        deterministic=deterministic,
        decode=config.decode)(x, attn_mask)

    x = nn.Dropout(rate=config.dropout_rate)(
        x, deterministic=deterministic)
    x = x + inputs
    # MLP block.
    y = nn.LayerNorm(dtype=config.dtype)(x)
    y = MlpBlock(config)(y, deterministic=deterministic)
    y = nn.Dropout(rate=config.dropout_rate)(
        y, deterministic=deterministic)

    return x + y
Applies Transformer module.

    Args:
      inputs: input data.
      deterministic: deterministic
      attn_mask: self-attention mask.

    Returns:
      output from VTNLayer: [batch_size, seq_length, model_dim].
    
videogvt/models/transformer_lm.py