class_docstr: Implements video-text fusion by self attention.
  
  We append the text CLS token to the video tokens and then feed the
  concatenated sequence into a Transformer.
  
  Attributes:
    self_attention_encoder_config: The config of the self attention encoder.
    self_attention_encoder_name: The type of self attention encoder.
docstr: Fuses video text by self attention.
  
  Args:
    video_tokens: A 3D float tensor of shape (batch_size, sequence_length,
      channels) representing the video tokens.
    text_embs: A 2D float tensor of shape (num_classes or batch_size *
      max_num_captions, channels) representing the text CLS token for each
      class. The first dimension is batch_size * max_num_captions if task =
      `moment_retrieval`. Otherwise, it is num_classes.
    task: 'action_segmentation', 'temporal_localization',
      'highlight_detection', or 'moment_retrieval'.
    input_word_ids: None.
    text_input_mask: None.
    video_input_mask: None or a 2D binary tensor of shape (batch_size,
      sequence_length).
    train: Whether or not it is in training.
  
  Returns:
    A 4D float tensor of shape (batch_size, batch_size * max_num_captions,
    sequence_length, channels) representing the pre_logits for each class if
    task = `moment_retrieval`.
  
    A 4D float tensor of shape (batch_size, num_classes, sequence_length,
    channels) representing the pre_logits for each class if task is
    `temporal_localization`.
  
    Video tokens in shape (batch_size, sequence_length, num_classes,
    channels) and text tokens in shape (batch_size, num_classes, channels) if
    task is `action_segmentation`.
code: def __call__(
      self,
      video_tokens: jnp.ndarray,
      text_embs: 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]:
      """Fuses video text by self attention.
  
      Args:
        video_tokens: A 3D float tensor of shape (batch_size, sequence_length,
          channels) representing the video tokens.
        text_embs: A 2D float tensor of shape (num_classes or batch_size *
          max_num_captions, channels) representing the text CLS token for each
          class. The first dimension is batch_size * max_num_captions if task =
          `moment_retrieval`. Otherwise, it is num_classes.
        task: 'action_segmentation', 'temporal_localization',
          'highlight_detection', or 'moment_retrieval'.
        input_word_ids: None.
        text_input_mask: None.
        video_input_mask: None or a 2D binary tensor of shape (batch_size,
          sequence_length).
        train: Whether or not it is in training.
  
      Returns:
        A 4D float tensor of shape (batch_size, batch_size * max_num_captions,
        sequence_length, channels) representing the pre_logits for each class if
        task = `moment_retrieval`.
  
        A 4D float tensor of shape (batch_size, num_classes, sequence_length,
        channels) representing the pre_logits for each class if task is
        `temporal_localization`.
  
        Video tokens in shape (batch_size, sequence_length, num_classes,
        channels) and text tokens in shape (batch_size, num_classes, channels) if
        task is `action_segmentation`.
      """
      encoder = _VIDEO_TEXT_ENCODER[self.self_attention_encoder_name](
          name="video_text_encoder",
          **self.self_attention_encoder_config,
      )
      if video_input_mask is None:
          feature_pyramid_config = (
              self.self_attention_encoder_config.get(
                  "feature_pyramid_config"
              )
          )
          if feature_pyramid_config is None:
              video_input_mask = jnp.ones(
                  video_tokens.shape[:2], dtype=jnp.int32
              )
          else:
              feature_pyramid_levels = (
                  feature_pyramid_config.feature_pyramid_levels
              )
              num_features_level0 = (
                  feature_pyramid_config.num_features_level0
              )
              feature_pyramid_downsample_stride = (
                  feature_pyramid_config.feature_pyramid_downsample_stride
              )
              fpn_video_token_len = sum(
                  [
                      num_features_level0
                      // (feature_pyramid_downsample_stride**idx)
                      for idx in range(len(feature_pyramid_levels))
                  ]
              )
              video_input_mask = jnp.ones(
                  (video_tokens.shape[0], fpn_video_token_len),
                  dtype=jnp.int32,
              )
      video_tokens, text_embs = jax.vmap(
          functools.partial(
              self._self_attention_encode_per_text_emb,
              encoder=encoder,
              task=task,
              train=train,
          ),
          in_axes=[None, None, 0],
      )(video_tokens, video_input_mask, text_embs)
      if task == "temporal_localization":
          # Converts 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_embs, [1, 0, 2]),
          )
      elif task == "action_segmentation":
          # Converts 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_embs, [1, 0, 2]),
          )
      elif task == "moment_retrieval":
          # Converts video_toekns 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_embs, [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 2 for calling the VideoTextEmbSelfAttentionFusion obj's underlying member method code with example usage: obj.__call__(video_tokens, text_embs, task, input_word_ids, text_input_mask, video_input_mask, train) and returns: jnp,  jnp, jnp,  jnp, jnp,  jnp