class_docstr: Transformer Model Decoder for sequence to sequence translation.
  
  Args:
    config: TransformerConfig dataclass containing hyperparameters.
    shared_embedding: a shared embedding layer to use.
class_docstr_tok: ["Transformer","Model","Decoder","for","sequence","to","sequence","translation",".","Args",":","config",":","TransformerConfig","dataclass","containing","hyperparameters",".","shared_embedding",":","a","shared","embedding","layer","to","use","."]
docstr: Applies Transformer model on the inputs.
  
  Args:
    encoded: encoded input data from encoder.
    inputs: input data.
    inputs_positions: input subsequence positions for packed examples.
    inputs_segmentation: input segmentation info for packed examples.
    decoder_mask: decoder self-attention mask.
    encoder_decoder_mask: encoder-decoder attention mask.
  
  Returns:
    output of a transformer decoder.
docstr_tok: ["Applies","Transformer","model","on","the","inputs",".","Args",":","encoded",":","encoded","input","data","from","encoder",".","inputs",":","input","data",".","inputs_positions",":","input","subsequence","positions","for","packed","examples",".","inputs_segmentation",":","input","segmentation","info","for","packed","examples",".","decoder_mask",":","decoder","self-attention","mask",".","encoder_decoder_mask",":","encoder-decoder","attention","mask",".","Returns",":","output","of","a","transformer","decoder","."]
code: def __call__(
      self,
      inputs,
      inputs_positions=None,
      inputs_segmentation=None,
      decoder_mask=None,
      encoder_decoder_mask=None,
  ):
      """Applies Transformer model on the inputs.
  
      Args:
        encoded: encoded input data from encoder.
        inputs: input data.
        inputs_positions: input subsequence positions for packed examples.
        inputs_segmentation: input segmentation info for packed examples.
        decoder_mask: decoder self-attention mask.
        encoder_decoder_mask: encoder-decoder attention mask.
  
      Returns:
        output of a transformer decoder.
      """
      config = self.config
      assert inputs.ndim == 2  # (batch, len)
  
      # Target Embedding
      if self.shared_embedding is None:
          output_embed = nn.Embed(
              num_embeddings=config.output_vocab_size,
              features=config.emb_dim,
              embedding_init=nn.initializers.normal(stddev=1.0),
          )
      else:
          output_embed = self.shared_embedding
  
      y = inputs.astype("int32")
      if not config.decode:
          y = shift_inputs(y, segment_ids=inputs_segmentation)
      y = output_embed(y)
      y = AddPositionEmbs(
          config=config, decode=config.decode, name="posembed_output"
      )(y, inputs_positions=inputs_positions)
      y = nn.Dropout(rate=config.dropout_rate)(
          y, deterministic=config.deterministic
      )
  
      y = y.astype(config.dtype)
  
      # Target-Input Decoder
      for lyr in range(config.num_layers):
          y = EncoderDecoder1DBlock(
              config=config, name=f"encoderdecoderblock_{lyr}"
          )(
              y,
              decoder_mask=decoder_mask,
              encoder_decoder_mask=encoder_decoder_mask,
          )
      y = nn.LayerNorm(dtype=config.dtype, name="encoderdecoder_norm")(
          y
      )
  
      # Decoded Logits
      if config.logits_via_embedding:
          # Use the transpose of embedding matrix for logit transform.
          logits = output_embed.attend(y.astype(jnp.float32))
          # Correctly normalize pre-softmax logits for this shared case.
          logits = logits / jnp.sqrt(y.shape[-1])
      else:
          logits = nn.Dense(
              config.output_vocab_size,
              dtype=config.dtype,
              kernel_init=config.kernel_init,
              bias_init=config.bias_init,
              name="logitdense",
          )(y)
      return logits
  
imports: from __future__ import annotations
  from typing import Callable, Any, Optional
  from flax import linen
  from flax import struct
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the Decoder obj's underlying member method code with example usage: obj.__call__(inputs, inputs_positions, inputs_segmentation, decoder_mask, encoder_decoder_mask) and returns: logits