def __call__(self, inputs, deterministic=True):
    """Applies Transformer decoder on the inputs.

    Args:
      inputs: input data
      deterministic: deterministic

    Returns:
      output: [batch_size, seq_length, model_dim]
    """
    config = self.config
    x = inputs.astype('int32')

    if config.decode:
      # decoder_mask will be defined in multiheadattention
      decoder_mask = None
    else:
      decoder_mask = nn.combine_masks(
          nn.make_attention_mask(inputs > 0, inputs > 0, dtype=config.dtype),
          nn.make_causal_mask(inputs, dtype=config.dtype))

    # Input Embedding
    input_embed = nn.Embed(
        num_embeddings=config.vocab_size,
        features=config.emb_dim,
        embedding_init=nn.initializers.normal(stddev=config.emb_dim**-0.5))

    x = input_embed(x)
    x = nn.Dropout(rate=config.input_dropout_rate)(
        x, deterministic=deterministic)
    x = PositionEmbeddings(config)(x)

    # Decoder
    for layer in range(config.num_layers):
      x = TransformerBlock(config, name=f'VTNLayer_{layer}')(
          x, deterministic=deterministic, attn_mask=decoder_mask)

    decoder_out = nn.LayerNorm(dtype=config.dtype, name='decoder_out_ln')(x)

    if config.logits_via_embedding:
      # Use the transpose of embedding matrix for logit transform.
      assert config.vocab_size == config.output_vocab_size
      logits = input_embed.attend(decoder_out.astype(jnp.float32))
      # Correctly normalize pre-softmax logits for this shared case.
      logits = logits / jnp.sqrt(decoder_out.shape[-1])
    else:
      logits = nn.Dense(
          config.output_vocab_size,
          dtype=config.dtype,
          kernel_init=nn.initializers.normal(
              stddev=decoder_out.shape[-1]**-0.5),
          bias_init=config.bias_init,
          name='logits_dense')(decoder_out)
    return logits
Applies Transformer decoder on the inputs.

    Args:
      inputs: input data
      deterministic: deterministic

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