class_docstr: Open-vocabulary instance class predictor.
class_docstr_tok: ["Open-vocabulary","instance","class","predictor","."]
docstr: Computes class prediction logits.
  
  Query embeddings from a text encoder define the classification label space.
  
  Args:
    x: Image features [batch_size, num_patches, emb_dim].
    query_embeddings: The embeddings to classify against of shape [batch_size,
      num_queries, out_dim]. If not specified, only the image class embeddings
      will be returned.
    query_mask: Mask indicating whether query is real (1) or padding (0), of
      shape [batch_size, num_queries].
  Returns:
    Dict with keys 'class_embeddings' and, if query embeddings were provided,
    'pred_logits'.
docstr_tok: ["Computes","class","prediction","logits",".","Query","embeddings","from","a","text","encoder","define","the","classification","label","space",".","Args",":","x",":","Image","features","[","batch_size",",","num_patches",",","emb_dim","]",".","query_embeddings",":","The","embeddings","to","classify","against","of","shape","[","batch_size",",","num_queries",",","out_dim","]",".","If","not","specified",",","only","the","image","class","embeddings","will","be","returned",".","query_mask",":","Mask","indicating","whether","query","is","real","(","1",")","or","padding","(","0",")",",","of","shape","[","batch_size",",","num_queries","]",".","Returns",":","Dict","with","keys","'class_embeddings","'","and",",","if","query","embeddings","were","provided",",","'pred_logits","'","."]
code: def __call__(
      self,
      x: jnp.ndarray,
      query_embeddings: Optional[jnp.ndarray] = None,
      query_mask: Optional[jnp.ndarray] = None,
  ) -> Dict[str, jnp.ndarray]:
      """Computes class prediction logits.
  
      Query embeddings from a text encoder define the classification label space.
  
      Args:
        x: Image features [batch_size, num_patches, emb_dim].
        query_embeddings: The embeddings to classify against of shape [batch_size,
          num_queries, out_dim]. If not specified, only the image class embeddings
          will be returned.
        query_mask: Mask indicating whether query is real (1) or padding (0), of
          shape [batch_size, num_queries].
      Returns:
        Dict with keys 'class_embeddings' and, if query embeddings were provided,
        'pred_logits'.
      """
      if self.out_dim is not None:
          out_dim = self.out_dim
      elif query_embeddings is not None:
          out_dim = query_embeddings.shape[-1]
      else:
          raise ValueError(
              "Unable to infer class head shape. Please pass out_dim."
          )
  
      image_class_emb = nn.Dense(
          out_dim, kernel_init=nn.initializers.normal(1e-6)
      )(x)
      if query_embeddings is None:
          return {"class_embeddings": image_class_emb}
      assert out_dim == query_embeddings.shape[-1]
  
      if self.normalize:
          image_class_emb /= (
              jnp.linalg.norm(image_class_emb, axis=-1, keepdims=True)
              + 1e-6
          )
          query_embeddings /= (
              jnp.linalg.norm(query_embeddings, axis=-1, keepdims=True)
              + 1e-6
          )
  
      assert query_embeddings.ndim > 2, (
          "Expects shape (batch, query, out_dim). "
          f"Got {query_embeddings.shape}"
      )
      pred_logits = jnp.einsum(
          "...pd,...qd->...pq", image_class_emb, query_embeddings
      )
  
      # Apply a learnable shift and scale to logits:
      logit_shift = nn.Dense(1, name="logit_shift")(x)
      logit_scale = nn.Dense(1, use_bias=True, name="logit_scale")(x)
      logit_scale = nn.elu(logit_scale) + 1
      pred_logits = (pred_logits + logit_shift) * logit_scale
  
      if query_mask is not None:
          if query_mask.ndim > 1:
              query_mask = jnp.expand_dims(query_mask, axis=-2)
          pred_logits = jnp.where(query_mask == 0, -1e6, pred_logits)
  
      return {
          "pred_logits": pred_logits,
          "class_embeddings": image_class_emb,
      }
  
imports: import abc
  import functools
  from typing import Any, Callable, Dict, Optional, Tuple, Union, Sequence
  from absl import logging
  from big_vision.models import bit
  from big_vision.models import vit
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import box_utils
  from scenic.projects.owl_vit import utils
  from scenic.projects.owl_vit.clip import layers
  from scenic.projects.owl_vit.clip import model
label_desc: Use image node_id 1 for calling the ClassPredictor obj's underlying member method code with example usage: obj.__call__(x, query_embeddings, query_mask) and returns: dict, dict