def cosine_similarity(a: Array, b: Array) -> Array:
  """Computes batched cosine similarity between two 2D arrays."""
  a_norm = jnp.linalg.norm(a, axis=-1)
  b_norm = jnp.linalg.norm(b, axis=-1)
  dot = jnp.einsum('bd,bd->b', a, b)
  return dot / (_SMALL + a_norm * b_norm)
Computes batched cosine similarity between two 2D arrays.
language/mentionmemory/utils/jax_utils.py