code: def test_video_text_emb_self_attention_fusion(
      self,
      task,
      video_input_mask,
      expected_video_shape,
      expected_text_shape,
  ):
      video_tokens = np.ones((2, 8, 16), dtype=np.float32)
      text_embs = (
          np.ones((2, 16), np.float32)
          if task == "moment_retrieval"
          else np.ones((4, 16), np.float32)
      )
      config = ml_collections.ConfigDict(
          {
              "num_heads": 2,
              "mlp_dim": 32,
              "num_layers": 2,
              "dropout_rate": 0.0,
              "attention_dropout_rate": 0.0,
              "stochastic_depth": 0.0,
          }
      )
      rng = random.PRNGKey(0)
      (
          actual_video_tokens,
          actual_text_tokens,
      ), _ = video_text_fusion.VideoTextEmbSelfAttentionFusion(
          self_attention_encoder_config=config
      ).init_with_output(
          rng,
          video_tokens,
          text_embs,
          video_input_mask=video_input_mask,
          task=task,
      )
      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 1 for calling the LayersTest obj's underlying member method code with example usage: obj.test_video_text_emb_self_attention_fusion(task, video_input_mask, expected_video_shape, expected_text_shape) without return types