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 all video-text pairs 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 3D float tensor of shape (batch_size, sequence_length,
      channels) representing the text tokens.
    input_word_ids: A 2D int tensor of shape (batch_size, sequence_length)
      representing the input word indices.
    text_input_mask: A 2D binary tensor of shape (batch_size, sequence_length)
      representing the mask of the text inputs.
    encoder: The Transformer encoder.
    train: Whether or not the model is under training.
  
  Returns:
    video_tokens: A 4D float tensor of shape (batch size, 1, sequence_length,
      channels) representing the fused frame embedding.
    text_emb: a 3D float tensor of shape (batch_size, 1, channels)
      representing the text CLS token.
code: def _self_attention_encode_all_video_text_pairs(
      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,
      train: bool,
  ) -> Tuple[jnp.ndarray, jnp.ndarray]:
      """Fuses all video-text pairs 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 3D float tensor of shape (batch_size, sequence_length,
          channels) representing the text tokens.
        input_word_ids: A 2D int tensor of shape (batch_size, sequence_length)
          representing the input word indices.
        text_input_mask: A 2D binary tensor of shape (batch_size, sequence_length)
          representing the mask of the text inputs.
        encoder: The Transformer encoder.
        train: Whether or not the model is under training.
  
      Returns:
        video_tokens: A 4D float tensor of shape (batch size, 1, sequence_length,
          channels) representing the fused frame embedding.
        text_emb: a 3D float tensor of shape (batch_size, 1, channels)
          representing the text CLS token.
      """
  
      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, text_tokens], axis=1)
          input_mask = model_utils.merge_pyramid_input_masks(
              video_input_masks, input_text_mask=text_input_mask
          )
          num_text_tokens = text_tokens.shape[1]
      else:
          tokens = jnp.concatenate(
              [
                  video_tokens,
                  model_utils.extract_emb(
                      text_tokens,
                      self.text_tower_classifier,
                      keepdims=True,
                      input_mask=text_input_mask,
                      input_word_ids=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=text_input_mask,
              input_word_ids=input_word_ids,
          )
      else:
          text_emb = tokens[:, -1]
      video_tokens = tokens[:, :-num_text_tokens]
      return (
          jnp.expand_dims(video_tokens, axis=1),
          jnp.expand_dims(text_emb, axis=1),
      )
  
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 1 for calling the VideoTextSelfAttentionFusion obj's underlying member method code with example usage: obj._self_attention_encode_all_video_text_pairs(video_tokens, video_input_mask, text_tokens, input_word_ids, text_input_mask, encoder, train) and returns: jnp,  jnp