class_docstr: Tests for ViT+ layers.
docstr: Tests image and text embedding with a CLIP model.
code: def test_clip_image_text_embedder(self):
      """Tests image and text embedding with a CLIP model."""
      batch_size, num_queries, seq_len, img_size, patch_size = (
          8,
          10,
          15,
          224,
          32,
      )
      side_patches = img_size // patch_size
      embed_configs = ml_collections.ConfigDict(
          dict(
              type="clip",
              variant="vit_b32",
              merge_class_token="drop",
              text_stochastic_droplayer_rate=0.1,
              vision_stochastic_droplayer_rate=0.1,
          )
      )
      images = jnp.ones((batch_size, img_size, img_size, 3))
      texts = jnp.ones(
          (batch_size, num_queries, seq_len), dtype=jnp.int32
      )
      rng = jax.random.PRNGKey(0)
  
      embedder = layers.ClipImageTextEmbedder(embed_configs)
  
      with self.subTest(name="images_and_text"):
          (img, txt), _ = embedder.init_with_output(
              rng, images=images, texts=texts
          )
          self.assertEqual(
              img.shape, (batch_size, side_patches**2, 768)
          )
          self.assertEqual(txt.shape, (batch_size, num_queries, 512))
  
      with self.subTest(name="only_images"):
          (img, _), _ = embedder.init_with_output(
              rng, images=images, texts=None
          )
          self.assertEqual(
              img.shape, (batch_size, side_patches**2, 768)
          )
  
      with self.subTest(name="only_text"):
          (_, txt), _ = embedder.init_with_output(
              rng, images=None, texts=texts
          )
          self.assertEqual(txt.shape, (batch_size, num_queries, 512))
  
file_path: scenic/scenic/projects/owl_vit/tests/layers_test.py
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import ml_collections
  from scenic.projects.owl_vit import layers
label_desc: Use image node_id 3 for calling the LayersTest obj's underlying member method code with example usage: obj.test_clip_image_text_embedder() without return types