class_docstr: Transformer Model for sequence to sequence translation.
class_docstr_tok: ["Transformer","Model","for","sequence","to","sequence","translation","."]
docstr: Applies Transformer encoder-branch on the inputs.
docstr_tok: ["Applies","Transformer","encoder-branch","on","the","inputs","."]
code: def encode(
      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],
      *,
      train: bool,
      debug: bool = False,
  ):
      """Applies Transformer encoder-branch on the inputs."""
      # TODO(valgab): Make attention masks for the case where input_segmentation
      # is not None
      x = {
          "rgb": x_rgb,
          "flow": x_flow,
          "spectrogram": x_spec,
          "wave": x_wave,
          "text": x_text,
      }
  
      encoded = self.encoder(x, train=train, debug=debug)
      encoding_dict = None
  
      if self.freeze_rgb_stream:
          encoded["rgb"] = jax.lax.stop_gradient(encoded["rgb"])
          encoded = jnp.concatenate(
              [encoded[m] for m in self.encoder_config.modality_fusion],
              axis=1,
          )
          logging.info("stop_gradient applied")
      elif self.add_masked_word_prediction_loss:
          encoding_dict = encoded
          encoded = jnp.concatenate(
              [encoded[m] for m in self.encoder_config.modality_fusion],
              axis=1,
          )
      return encoded, encoding_dict
  
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 2 for calling the Seq2SeqModule obj's underlying member method code with example usage: obj.encode(x_rgb, x_flow, x_spec, x_wave, x_text) and returns: encoded,  encoding_dict