class_docstr: Implements video-text fusion by self attention.
  
  We concatenate all text tokens (or only the CLS token) with video tokens and
  then feed the concatenated sequence into a Transformer.
  
  Attributes:
    text_tower_classifier: 'token' (take the first token), 'eos' (take the last
    token).
    self_attention_encoder_config: The config of the self attention encoder.
    use_all_text_tokens: Whether or not to fuse with all text tokens. If False,
      we only fuse with the text CLS token.
docstr: Executes video-text fusion by self attention.
  
  Args:
    video_tokens: A 3D float tensor of shape (batch_size, sequence_length,
      channels) representing the video tokens.
    text_tokens: A 3D float tensor of shape (num_texts, sequence_length,
      channels) representing the text tokens. The first dimension is
      batch_size * max_num_captions if task = `moment_retrieval`. Otherwise,
      it is num_classes.
    task: 'action_segmentation', 'temporal_localization', 'moment_retrieval'
      or 'highlight_detection.
    input_word_ids: A 2D int tensor of shape (num_texts, sequence_length)
      representing the input word indices. This arg is used to find EOS id.
    text_input_mask: A 2D binary tensor of shape (num_texts, sequence_length)
      representing the mask of the text inputs.
    video_input_mask: A 2D binary tensor of shape (batch_size,
      sequence_length) representing the mask of the video inputs.
    train: Whether or not it is in training.
  
  Returns:
    video_tokens:
      A 4D float tensor of shape (batch_size, num_classes, sequence_length,
      channels) representing the fused tokens for each class if task is
      'temporal_localization'.
  
      A 4D float tensor of shape (batch_size, 1, sequence_length, channels)
      representing the fused frame tokens if task is 'highlight_detection'.
  
      A 3D float tensor of shape (batch_size, sequence_length, num_classes)
      representing the fused tokens for each class if task is
      'action_segmentation'.
  
      A 4D float tensor of shape (batch_size, batch_size * max_num_captions,
      sequence_length, channels) representing the fused tokens for each class
      if task = 'moment_retrieval'.
    text_tokens:
      Not used. A 3D float tensor of shape (batch_size, num_texts, channels)
      representing the text CLS token for each class. The second dimension is
      batch_size * max_num_captions if task = `moment_retrieval`. Otherwise,
      it is num_classes.
code: def __call__(
      self,
      video_tokens: jnp.ndarray,
      text_tokens: jnp.ndarray,
      task: str,
      input_word_ids: Optional[jnp.ndarray] = None,
      text_input_mask: Optional[jnp.ndarray] = None,
      video_input_mask: Optional[jnp.ndarray] = None,
      train: bool = False,
  ) -> Tuple[jnp.ndarray, jnp.ndarray]:
      """Executes video-text fusion by self attention.
  
      Args:
        video_tokens: A 3D float tensor of shape (batch_size, sequence_length,
          channels) representing the video tokens.
        text_tokens: A 3D float tensor of shape (num_texts, sequence_length,
          channels) representing the text tokens. The first dimension is
          batch_size * max_num_captions if task = `moment_retrieval`. Otherwise,
          it is num_classes.
        task: 'action_segmentation', 'temporal_localization', 'moment_retrieval'
          or 'highlight_detection.
        input_word_ids: A 2D int tensor of shape (num_texts, sequence_length)
          representing the input word indices. This arg is used to find EOS id.
        text_input_mask: A 2D binary tensor of shape (num_texts, sequence_length)
          representing the mask of the text inputs.
        video_input_mask: A 2D binary tensor of shape (batch_size,
          sequence_length) representing the mask of the video inputs.
        train: Whether or not it is in training.
  
      Returns:
        video_tokens:
          A 4D float tensor of shape (batch_size, num_classes, sequence_length,
          channels) representing the fused tokens for each class if task is
          'temporal_localization'.
  
          A 4D float tensor of shape (batch_size, 1, sequence_length, channels)
          representing the fused frame tokens if task is 'highlight_detection'.
  
          A 3D float tensor of shape (batch_size, sequence_length, num_classes)
          representing the fused tokens for each class if task is
          'action_segmentation'.
  
          A 4D float tensor of shape (batch_size, batch_size * max_num_captions,
          sequence_length, channels) representing the fused tokens for each class
          if task = 'moment_retrieval'.
        text_tokens:
          Not used. A 3D float tensor of shape (batch_size, num_texts, channels)
          representing the text CLS token for each class. The second dimension is
          batch_size * max_num_captions if task = `moment_retrieval`. Otherwise,
          it is num_classes.
      """
  
      encoder = _VIDEO_TEXT_ENCODER[self.self_attention_encoder_name](
          name="video_text_encoder",
          **self.self_attention_encoder_config,
      )
      if video_input_mask is None:
          video_input_mask = jnp.ones(
              video_tokens.shape[:2], dtype=jnp.int32
          )
      if task == "highlight_detection":
          return self._self_attention_encode_all_video_text_pairs(
              video_tokens,
              video_input_mask,
              text_tokens,
              input_word_ids,
              text_input_mask,
              encoder,
              train,
          )
  
      video_tokens, text_tokens = jax.vmap(
          functools.partial(
              self._self_attention_encode_per_text,
              encoder=encoder,
              task=task,
              train=train,
          ),
          in_axes=[None, None, 0, 0, 0],
      )(
          video_tokens,
          video_input_mask,
          text_tokens,
          input_word_ids,
          text_input_mask,
      )
      if task == "temporal_localization":
          # Convert video_tokens from (num_classes, batch_size, num_frames,
          # channels) to (batch_size, num_classes, num_frames, channels).
          return (
              jnp.transpose(video_tokens, [1, 0, 2, 3]),
              jnp.transpose(text_tokens, [1, 0, 2]),
          )
      elif task == "action_segmentation":
          # Convert video_tokens from (num_classes, batch_size, num_frames,
          # channels) to (batch_size, num_frames, num_classes, channels).
          return (
              jnp.transpose(video_tokens, [1, 2, 0, 3]),
              jnp.transpose(text_tokens, [1, 0, 2]),
          )
      elif task == "moment_retrieval":
          # Convert video_tokens from (batch_size * max_num_captions, batch_size,
          # num_frames, channels) to (batch_size, batch_size * max_num_captions,
          # num_frames, channels).
          return (
              jnp.transpose(video_tokens, [1, 0, 2, 3]),
              jnp.transpose(text_tokens, [1, 0, 2]),
          )
      else:
          raise ValueError(f"Unexpected task `{task}`.")
  
imports: import functools
  from typing import Callable, List, Optional, Sequence, Tuple
  from absl import logging
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.layers import attention_layers
  from scenic.projects.baselines import vit
  from scenic.projects.unloc import encoders
  from scenic.projects.unloc import model_utils
label_desc: Use image node_id 3 for calling the VideoTextSelfAttentionFusion obj's underlying member method code with example usage: obj.__call__(video_tokens, text_tokens, task, input_word_ids, text_input_mask, video_input_mask, train) and returns: self, jnp,  jnp, jnp,  jnp, jnp,  jnp