class_docstr: Transformer Model for sequence to sequence translation.
  
  Attributes:
    config: TransformerConfig dataclass containing hyperparameters.
class_docstr_tok: ["Transformer","Model","for","sequence","to","sequence","translation",".","Attributes",":","config",":","TransformerConfig","dataclass","containing","hyperparameters","."]
docstr: Applies Transformer decoder-branch on encoded-input and target.
  
  Args:
    encoded: encoded input data from encoder.
    inputs: input data (only needed for masking).
    targets: target data.
    targets_positions: target subsequence positions for packed examples.
    inputs_segmentation: input segmentation info for packed examples.
    targets_segmentation: target segmentation info for packed examples.
  
  Returns:
    logits array from transformer decoder.
docstr_tok: ["Applies","Transformer","decoder-branch","on","encoded-input","and","target",".","Args",":","encoded",":","encoded","input","data","from","encoder",".","inputs",":","input","data","(","only","needed","for","masking",")",".","targets",":","target","data",".","targets_positions",":","target","subsequence","positions","for","packed","examples",".","inputs_segmentation",":","input","segmentation","info","for","packed","examples",".","targets_segmentation",":","target","segmentation","info","for","packed","examples",".","Returns",":","logits","array","from","transformer","decoder","."]
code: def decode(
      self,
      encoded,
      inputs,  # only needed for masks
      targets,
      targets_positions=None,
      inputs_segmentation=None,
      targets_segmentation=None,
  ):
      """Applies Transformer decoder-branch on encoded-input and target.
  
      Args:
        encoded: encoded input data from encoder.
        inputs: input data (only needed for masking).
        targets: target data.
        targets_positions: target subsequence positions for packed examples.
        inputs_segmentation: input segmentation info for packed examples.
        targets_segmentation: target segmentation info for packed examples.
  
      Returns:
        logits array from transformer decoder.
      """
      config = self.config
  
      # Make padding attention masks.
      if config.decode:
          # for fast autoregressive decoding only a special encoder-decoder mask is used
          decoder_mask = None
          encoder_decoder_mask = nn.make_attention_mask(
              jnp.ones_like(targets) > 0, inputs > 0, dtype=config.dtype
          )
      else:
          decoder_mask = nn.combine_masks(
              nn.make_attention_mask(
                  targets > 0, targets > 0, dtype=config.dtype
              ),
              nn.make_causal_mask(targets, dtype=config.dtype),
          )
          encoder_decoder_mask = nn.make_attention_mask(
              targets > 0, inputs > 0, dtype=config.dtype
          )
  
      # Add segmentation block-diagonal attention masks if using segmented data.
      if inputs_segmentation is not None:
          decoder_mask = nn.combine_masks(
              decoder_mask,
              nn.make_attention_mask(
                  targets_segmentation,
                  targets_segmentation,
                  jnp.equal,
                  dtype=config.dtype,
              ),
          )
          encoder_decoder_mask = nn.combine_masks(
              encoder_decoder_mask,
              nn.make_attention_mask(
                  targets_segmentation,
                  inputs_segmentation,
                  jnp.equal,
                  dtype=config.dtype,
              ),
          )
      logits = self.decoder(
          encoded,
          targets,
          targets_positions=targets_positions,
          decoder_mask=decoder_mask,
          encoder_decoder_mask=encoder_decoder_mask,
      )
      return logits.astype(self.config.dtype)
  
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 3 for calling the Transformer obj's underlying member method code with example usage: obj.decode(encoded, inputs, targets, targets_positions, inputs_segmentation, targets_segmentation) and returns: logits