class_docstr: Transformer Model for sequence to sequence translation.
class_docstr_tok: ["Transformer","Model","for","sequence","to","sequence","translation","."]
docstr: Applies Transformer model on the inputs.
docstr_tok: ["Applies","Transformer","model","on","the","inputs","."]
code: def __call__(
      self,
      x_rgb: Optional[jnp.ndarray],
      x_flow: Optional[jnp.ndarray],
      x_spec: Optional[jnp.ndarray],
      x_wave: Optional[jnp.ndarray],
      x_text: Optional[jnp.ndarray],
      targets,
      masked_token_idxs: Optional[jnp.ndarray] = None,
      masked_token_idx_masks: Optional[jnp.ndarray] = None,
      masked_word_targets: Optional[jnp.ndarray] = None,
      decode: bool = False,
      *,
      train: bool,
      debug: bool = False,
  ):
      """Applies Transformer model on the inputs."""
  
      encoded = self.encode(
          x_rgb,
          x_flow,
          x_spec,
          x_wave,
          x_text,
          train=train,
          debug=debug,
      )
  
      output = self.decode(
          encoded[0], targets, decode=decode, train=train
      )
  
      if not train or not self.add_masked_word_prediction_loss:
          return output
  
      assert masked_token_idxs is not None
      assert masked_token_idx_masks is not None
      assert masked_word_targets is not None
      assert encoded[1] is not None
      logging.info("encoded[0] %s", encoded[0])
      logging.info("encoded[1] %s", encoded[1])
      max_num_masked_words = masked_token_idxs.shape[1]
      x_out = []
      x_mask = []
      sample_masked_inputs = jax.vmap(
          jax.vmap(lambda x, y: x[y], (None, 0), 0), (0, 0), 0
      )
      for modality in self.encoder_config.modality_fusion:
          modality_feature = encoded[1][modality]
          if modality == "spectrogram":
              logging.info("spectrogram feature %s", modality_feature)
              modality_feature_mask = masked_token_idx_masks
              if self.encoder_config.classifier == "token":
                  cls_token = sample_masked_inputs(
                      modality_feature,
                      jnp.zeros_like(masked_word_targets[..., 0:1]),
                  )
                  modality_feature = modality_feature[:, 1:, :]
                  cls_token_mask = jnp.ones_like(
                      masked_token_idx_masks[..., 0:1]
                  )
              modality_feature = sample_masked_inputs(
                  modality_feature, masked_token_idxs
              )
              if self.encoder_config.classifier == "token":
                  modality_feature = jnp.concatenate(
                      [cls_token, modality_feature], 2
                  )
                  modality_feature_mask = jnp.concatenate(
                      [cls_token_mask, masked_token_idx_masks], 2
                  )
              logging.info("spectrogram feature 2 %s", modality_feature)
          else:
              modality_feature = jnp.repeat(
                  modality_feature[:, jnp.newaxis],
                  max_num_masked_words,
                  1,
              )
              modality_feature_mask = jnp.ones_like(
                  modality_feature[..., 0]
              )
          x_out.append(modality_feature)
          x_mask.append(modality_feature_mask)
      masked_input_features = jnp.concatenate(x_out, 2)
      masked_input_feature_masks = jnp.concatenate(x_mask, 2)
      logging.info("masked_input_features %s", masked_input_features)
      b, m, t, e = masked_input_features.shape
      masked_input_features = jnp.reshape(
          masked_input_features, [b * m, t, e]
      )
      masked_input_masks = jnp.reshape(
          masked_input_feature_masks, [b * m, t]
      )
      masked_word_targets = jnp.reshape(
          masked_word_targets, [b * m, -1]
      )
  
      word_pred_output = self.decode(
          masked_input_features,
          masked_word_targets,
          decode=False,
          train=False,
          encoded_mask=masked_input_masks,
      )
  
      return output, word_pred_output
  
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 4 for calling the Seq2SeqModule obj's underlying member method code with example usage: obj.__call__(x_rgb, x_flow, x_spec, x_wave, x_text, targets, masked_token_idxs, masked_token_idx_masks, masked_word_targets, decode) and returns: output,  word_pred_output, output