def cosine_similarity(a, b):
  """Compute cosine similarity along final dimension of tensors.

  In tensor size descriptions, ... indicates an arbitrary number of dimensions.
  For example, this can compute similarity for any of these:
  - Dialogs: (batch_size, n_messages, n_words, size)
  - Text: (batch_size, n_words, size)
  - Hidden State: (batch_size, size)

  Args:
    a: Left tensor to compare, (batch_size, ..., n)
    b: Right tensor to compare (batch_size, ..., n)

  Returns:
    TF Scalar cosine similarity between vectors a and b
  """
  # (batch_size, ..., 1)
  a_norm = tf.norm(a, 2, axis=-1)
  # (batch_size, ..., 1)
  b_norm = tf.norm(b, 2, axis=-1)
  # (batch_size, ..., 1)
  return tf.reduce_sum(a * b, axis=-1) / (a_norm * b_norm)
Compute cosine similarity along final dimension of tensors.

  In tensor size descriptions, ... indicates an arbitrary number of dimensions.
  For example, this can compute similarity for any of these:
  - Dialogs: (batch_size, n_messages, n_words, size)
  - Text: (batch_size, n_words, size)
  - Hidden State: (batch_size, size)

  Args:
    a: Left tensor to compare, (batch_size, ..., n)
    b: Right tensor to compare (batch_size, ..., n)

  Returns:
    TF Scalar cosine similarity between vectors a and b
  
language/serene/layers.py