class_docstr: Wraps an OWL-ViT FLAX model for convenient inference.
  
  All public methods apply to a single example and take and return Numpy arrays.
  
  Attributes:
    config: ConfigDict with model configuration.
    module: OWL-ViT Flax module.
    variables: Variable dict to be used with module.apply.
docstr: Embeds image and returns image features, class embeddings, and boxes.
docstr_tok: ["Embeds","image","and","returns","image","features",",","class","embeddings",",","and","boxes","."]
code: def _embed_image_jitted(
      self, image: jnp.ndarray
  ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
      """Embeds image and returns image features, class embeddings, and boxes."""
      feature_map = self.module.apply(
          self.variables,
          images=image,
          train=False,
          method=self.module.image_embedder,
      )
      b, _, _, c = feature_map.shape
      features = jnp.reshape(feature_map, (b, -1, c))
      pred_boxes = self.module.apply(
          self.variables,
          image_features=features,
          feature_map=feature_map,
          method=self.module.box_predictor,
      )["pred_boxes"]
      class_embeddings = self.module.apply(
          self.variables,
          image_features=features,
          query_embeddings=None,
          method=self.module.class_predictor,
      )["class_embeddings"]
      return features, class_embeddings, pred_boxes
  
imports: import dataclasses
  import functools
  from typing import Any, Dict, Tuple
  from flax import 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.notebooks import numpy_cache
  from scipy import special
  from skimage import transform
  import tensorflow
label_desc: Use image node_id 9 for calling the Model obj's underlying member method code with example usage: obj._embed_image_jitted(image) and returns: features,  class_embeddings,  pred_boxes