class_docstr: Tests for ViT+ layers.
docstr: Tests class predictor.
code: def test_class_predictor(self):
      """Tests class predictor."""
      batch_size, num_queries, img_size, patch_size = 8, 10, 224, 32
      side_patches = img_size // patch_size
      num_patches = side_patches**2
      img_emb_dim, out_dim = 64, 100
      img_emb = jnp.ones((batch_size, num_patches, img_emb_dim))
      txt_emb = jnp.ones((batch_size, num_queries, out_dim))
      label_mask = jnp.ones((batch_size, num_queries), dtype=jnp.int32)
  
      class_predictor = layers.ClassPredictor(out_dim=out_dim)
      outputs, variables = class_predictor.init_with_output(
          jax.random.PRNGKey(0), img_emb, txt_emb, label_mask
      )
      self.assertEqual(
          outputs["pred_logits"].shape,
          (batch_size, num_patches, num_queries),
      )
      num_weights = jax.tree_util.tree_flatten(variables)[0]
      self.assertLen(num_weights, 6)
  
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 2 for calling the LayersTest obj's underlying member method code with example usage: obj.test_class_predictor() without return types