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: Fuses video and text tokens per class via self attention.
  
  Args:
    video_tokens: A 3D float tensor of shape (batch_size, sequence_length,
      channels) representing the video tokens.
    video_input_mask: A 2D binary tensor of shape (batch_size,
      sequence_length).
    text_tokens: A 2D float tensor of shape (sequence_length, channels)
      representing the text tokens.
    input_word_ids: A 1D int tensor of shape (sequence_length) representing
      the input word indices.
    text_input_mask: A 1D binary tensor of shape (sequence_length)
      representing the mask of the text inputs.
    encoder: The Transformer encoder.
    task: 'action_segmentation', 'moment_retrieval', 'highlight_detection', or
      'temporal_localization'.
    train: Whether or not the model is under training.
  
  Returns:
    video_tokens: A 3D float tensor of shape (batch size, sequence_length,
      channels) representing the fused frame embedding.
    text_emb: A 2D float tensor of shape (batch_size, channels) representing
      the text CLS token.
code: def _self_attention_encode_per_text(
      self,
      video_tokens: jnp.ndarray,
      video_input_mask: jnp.ndarray,
      text_tokens: jnp.ndarray,
      input_word_ids: jnp.ndarray,
      text_input_mask: jnp.ndarray,
      encoder: nn.Module,
      task: str,
      train: bool,
  ) -> Tuple[jnp.ndarray, jnp.ndarray]:
      """Fuses video and text tokens per class via self attention.
  
      Args:
        video_tokens: A 3D float tensor of shape (batch_size, sequence_length,
          channels) representing the video tokens.
        video_input_mask: A 2D binary tensor of shape (batch_size,
          sequence_length).
        text_tokens: A 2D float tensor of shape (sequence_length, channels)
          representing the text tokens.
        input_word_ids: A 1D int tensor of shape (sequence_length) representing
          the input word indices.
        text_input_mask: A 1D binary tensor of shape (sequence_length)
          representing the mask of the text inputs.
        encoder: The Transformer encoder.
        task: 'action_segmentation', 'moment_retrieval', 'highlight_detection', or
          'temporal_localization'.
        train: Whether or not the model is under training.
  
      Returns:
        video_tokens: A 3D float tensor of shape (batch size, sequence_length,
          channels) representing the fused frame embedding.
        text_emb: A 2D float tensor of shape (batch_size, channels) representing
          the text CLS token.
      """
      tiled_text_tokens = jnp.tile(
          jnp.expand_dims(text_tokens, axis=0),
          [video_tokens.shape[0], 1, 1],
      )
      tiled_text_mask = jnp.tile(
          jnp.expand_dims(text_input_mask, axis=0),
          [video_tokens.shape[0], 1],
      )
      tiled_input_word_ids = jnp.tile(
          jnp.expand_dims(input_word_ids, axis=0),
          [video_tokens.shape[0], 1],
      )
      feature_pyramid_config = self.self_attention_encoder_config.get(
          "feature_pyramid_config"
      )
      if feature_pyramid_config is None:
          video_input_masks = [video_input_mask]
      else:
          video_input_masks = model_utils.create_pyramid_input_masks(
              video_input_mask,
              num_features_level0=feature_pyramid_config.num_features_level0,
              num_pyramid_levels=len(
                  feature_pyramid_config.feature_pyramid_levels
              ),
              feature_pyramid_downsample_stride=feature_pyramid_config.feature_pyramid_downsample_stride,
              num_text_tokens=0,
          )
      if self.use_all_text_tokens:
          tokens = jnp.concatenate(
              [video_tokens, tiled_text_tokens], axis=1
          )
          input_mask = model_utils.merge_pyramid_input_masks(
              video_input_masks, input_text_mask=tiled_text_mask
          )
          num_text_tokens = text_tokens.shape[0]
      else:
          tokens = jnp.concatenate(
              [
                  video_tokens,
                  model_utils.extract_emb(
                      tiled_text_tokens,
                      self.text_tower_classifier,
                      keepdims=True,
                      input_mask=tiled_text_mask,
                      input_word_ids=tiled_input_word_ids,
                  ),
              ],
              axis=1,
          )
          input_mask = model_utils.merge_pyramid_input_masks(
              video_input_masks,
              input_text_mask=jnp.ones(
                  (video_tokens.shape[0], 1), dtype=jnp.int32
              ),
          )
          num_text_tokens = 1
  
      tokens = encoder(tokens, input_mask=input_mask, train=train)
      if self.use_all_text_tokens:
          text_emb = model_utils.extract_emb(
              tokens[:, -num_text_tokens:],
              self.text_tower_classifier,
              keepdims=False,
              input_mask=tiled_text_mask,
              input_word_ids=tiled_input_word_ids,
          )
      else:
          text_emb = tokens[:, -1]
      if task in {
          "action_segmentation",
          "temporal_localization",
          "moment_retrieval",
      }:
          return tokens[:, :-num_text_tokens], text_emb
      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 VideoTextSelfAttentionFusion obj's underlying member method code with example usage: obj._self_attention_encode_per_text(video_tokens, video_input_mask, text_tokens, input_word_ids, text_input_mask, encoder, task, train) and returns: tokens,  text_emb