class_docstr: 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.
  
  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",".","Returns",":","output","of","a","transformer","decoder","."]
code: def __call__(
      self,
      inputs,
      *,
      inputs_positions=None,
      inputs_segmentation=None,
      decoder_mask=None,
      rngs: nnx.Rngs | None = 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.
  
      Returns:
        output of a transformer decoder.
      """
      config = self.config
      assert inputs.ndim == 2  # (batch, len)
  
      y = inputs.astype("int32")
      if not flaglib.flags.get("decode", False):
          y = shift_inputs(y, segment_ids=inputs_segmentation)
      y = self.output_embed(y)
      y = self.posembed_output(y, inputs_positions=inputs_positions)
      y = self.dropout(y, rngs=rngs)
  
      y = y.astype(config.dtype)
  
      # Target-Input Decoder
      for idx in range(config.num_layers):
          # TODO(cgarciae): use a list of layers instead of getattr
          layer: EncoderDecoder1DBlock = getattr(
              self, f"encoderdecoderblock_{idx}"
          )
          y = layer(
              y,
              decoder_mask=decoder_mask,
              rngs=rngs,
          )
      y = self.encoderdecoder_norm(y)
  
      # Decoded Logits
      if self.logitdense:
          logits = self.logitdense(y)
      else:
          # Use the transpose of embedding matrix for logit transform.
          logits = self.output_embed.attend(y.astype(jnp.float32))
          # Correctly normalize pre-softmax logits for this shared case.
          logits = logits / jnp.sqrt(y.shape[-1])
      return logits
  
imports: from __future__ import annotations
  import dataclasses
  from typing import Any, Optional
  import jax
  import jax.numpy
  import numpy
  from jax import lax
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import flaglib
  from flax.experimental.nnx.examples.lm1b.configs import default
label_desc: Use image node_id 2 for calling the Decoder obj's underlying member method code with example usage: obj.__call__(inputs) and returns: logits