class_docstr: Transformer Model for sequence to sequence translation.
class_docstr_tok: ["Transformer","Model","for","sequence","to","sequence","translation","."]
docstr: Applies Transformer decoder-branch on encoded-input and target.
  
  Args:
    encoded: encoded input data from encoder.
    targets: target data.
    decode: whether to run in single-position autoregressive mode.
    train: whether to apply dropout
    encoded_mask: mask tensor indicating valitity of each token in encoded.
    debug: debug mode
  
  Returns:
    logits array from transformer decoder.
docstr_tok: ["Applies","Transformer","decoder-branch","on","encoded-input","and","target",".","Args",":","encoded",":","encoded","input","data","from","encoder",".","targets",":","target","data",".","decode",":","whether","to","run","in","single-position","autoregressive","mode",".","train",":","whether","to","apply","dropout","encoded_mask",":","mask","tensor","indicating","valitity","of","each","token","in","encoded",".","debug",":","debug","mode","Returns",":","logits","array","from","transformer","decoder","."]
code: def decode(
      self,
      encoded,
      targets,  # Used for teacher forcing
      decode: bool,
      train: bool,
      encoded_mask: Optional[jnp.ndarray] = None,
      debug: bool = False,
  ):
      """Applies Transformer decoder-branch on encoded-input and target.
  
      Args:
        encoded: encoded input data from encoder.
        targets: target data.
        decode: whether to run in single-position autoregressive mode.
        train: whether to apply dropout
        encoded_mask: mask tensor indicating valitity of each token in encoded.
        debug: debug mode
  
      Returns:
        logits array from transformer decoder.
      """
      cfg = self.decoder_config
  
      # Make padding attention masks.
      if decode:
          # for fast autoregressive decoding only a special encoder-decoder mask is
          # used.
          decoder_mask = None
      else:
          # Teacher forcing
          # No attention to target paddings, no attention to future tokens
          decoder_mask = nn.combine_masks(
              nn.make_attention_mask(
                  targets > 0, targets > 0, dtype=cfg.dtype
              ),
              nn.make_causal_mask(targets, dtype=self.dtype),
          )
      encoder_decoder_mask = None
      if encoded_mask is not None:
          encoder_decoder_mask = encoded_mask[
              :, jnp.newaxis, jnp.newaxis, :
          ]
      logits = self.decoder(
          encoded,
          targets,
          decoder_mask=decoder_mask,
          encoder_decoder_mask=encoder_decoder_mask,
          decode=decode,
          train=train,
      )
      return logits.astype(self.dtype)
  
imports: from typing import Any, Dict, Optional, Tuple
  from absl import logging
  from flax import linen
  from flax.training import common_utils
  from immutabledict import immutabledict
  import jax
  from jax import lax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.mbt import model
  from scenic.projects.mbt.model import temporal_encode
label_desc: Use image node_id 3 for calling the Seq2SeqModule obj's underlying member method code with example usage: obj.decode(encoded, targets, decode, train, encoded_mask, debug) and returns: logits