def topk_search_fn(query_embs, embeddings_mat,
                   k):
  """Compute topk of inner product search.

  Args:
    query_embs: batch_size, hidden_size
    embeddings_mat: num_vocab, hidden_size
    k: top k

  Returns:
    topk_fact_logits, topk_fact_ids
  """
  fact_logits = tf.matmul(
      query_embs, embeddings_mat, transpose_b=True)
  # batch_size, num_facts
  fact_logits /= tf.sqrt(tf.cast(tf.shape(query_embs)[1], tf.float32))
  # batch_size, num_facts
  topk_fact_logits, topk_fact_ids = tf.nn.top_k(fact_logits, k=k)
  # batch_size, intermediate_top_k
  return topk_fact_logits, topk_fact_ids
Compute topk of inner product search.

  Args:
    query_embs: batch_size, hidden_size
    embeddings_mat: num_vocab, hidden_size
    k: top k

  Returns:
    topk_fact_logits, topk_fact_ids
  
language/emql/util.py