class_docstr: A two-tower image-text encoder model.
class_docstr_tok: ["A","two-tower","image-text","encoder","model","."]
docstr: Computes the similarity score matrix from `a` to `b`.
docstr_tok: ["Computes","the","similarity","score","matrix","from","`","a","`","to","`","b","`","."]
code: def compute_similarity(
      self,
      a: jnp.ndarray,
      b: jnp.ndarray,
      all_gather_axis_name: Optional[str] = None,
  ) -> jnp.ndarray:
      """Computes the similarity score matrix from `a` to `b`."""
      assert self.similarity in {"cosine", "dot_product"}
  
      if self.similarity == "cosine":
          a /= (
              jnp.linalg.norm(a, axis=-1, keepdims=True)
              + jnp.finfo(a.dtype).eps
          )
          b /= (
              jnp.linalg.norm(b, axis=-1, keepdims=True)
              + jnp.finfo(b.dtype).eps
          )
  
      if all_gather_axis_name:
          a = jax.lax.all_gather(a, all_gather_axis_name).reshape(
              -1, a.shape[-1]
          )
          b = jax.lax.all_gather(b, all_gather_axis_name).reshape(
              -1, b.shape[-1]
          )
  
      return jnp.exp(self.logit_scale) * a @ b.T
  
imports: import abc
  from typing import Literal
  from typing import Optional
  import flax
  from flax.core.frozen_dict import FrozenDict
  import flax.linen
  import jax
  import jax.numpy
label_desc: Use image node_id 8 for calling the ImageTextEncoder obj's underlying member method code with example usage: obj.compute_similarity(a, b, all_gather_axis_name) and returns: unknown