code: def test_video_text_self_attention_fusion_fpn(
      self, task, expected_video_shape, expected_text_shape
  ):
      batch_size = 2
      num_classes = 4
      video_seq = 8
      text_seq = 10
      video_tokens = np.ones(
          (batch_size, video_seq, 16), dtype=np.float32
      )
      if task == "moment_retrieval" or task == "highlight_detection":
          text_tokens = np.ones((batch_size, text_seq, 16), np.float32)
          input_word_ids = np.ones((batch_size, text_seq), np.int32)
          text_input_mask = np.ones((batch_size, text_seq), np.int32)
      else:
          text_tokens = np.ones((num_classes, text_seq, 16), np.float32)
          input_word_ids = np.ones((num_classes, text_seq), np.int32)
          text_input_mask = np.ones((num_classes, text_seq), np.int32)
      video_input_mask = np.ones(
          (batch_size, video_seq + video_seq // 2), np.int32
      )
      config = ml_collections.ConfigDict(
          {
              "num_heads": 2,
              "mlp_dim": 32,
              "num_layers": 3,
              "dropout_rate": 0.0,
              "attention_dropout_rate": 0.0,
              "stochastic_depth": 0.0,
              "feature_pyramid_config": ml_collections.ConfigDict(
                  {
                      "num_features_level0": 8,
                      "feature_pyramid_levels": [1, 2],
                      "feature_pyramid_downsample_stride": 2,
                  }
              ),
          }
      )
      rng = random.PRNGKey(0)
      (
          actual_video_tokens,
          actual_text_tokens,
      ), _ = video_text_fusion.VideoTextSelfAttentionFusion(
          text_tower_classifier="eos",
          self_attention_encoder_name="fpn",
          self_attention_encoder_config=config,
          use_all_text_tokens=False,
      ).init_with_output(
          rng,
          video_tokens,
          text_tokens,
          task,
          input_word_ids,
          text_input_mask,
          video_input_mask,
      )
      self.assertTupleEqual(
          actual_video_tokens.shape, expected_video_shape
      )
      self.assertTupleEqual(
          actual_text_tokens.shape, expected_text_shape
      )
  
file_path: scenic/scenic/projects/unloc/video_text_fusion_test.py
imports: from absl.testing import parameterized
  from jax import random
  import ml_collections
  import numpy
  from scenic.projects.unloc import video_text_fusion
  import tensorflow
label_desc: Use image node_id 4 for calling the LayersTest obj's underlying member method code with example usage: obj.test_video_text_self_attention_fusion_fpn(task, expected_video_shape, expected_text_shape) without return types