def compute_query_embs_and_sketch(
      self, entity_ids, logits,
      params):
    """Compute embeddings and sketch with logits of entities.

    Given entity_ids and logits in the same order, compute weighted average
    of entity embeddings and a sketch that stores their weights.

    Args:
      entity_ids: entity ids
      logits: logits before softmax
      params: params

    Returns:
      A weighted average of embeddings and a cm sketch for the entities
    """
    topk_softmax = tf.nn.softmax(logits)
    # batch_size, topk
    topk_ent_embs = tf.nn.embedding_lookup(self.entity_embeddings_mat,
                                           entity_ids)
    # batch_size, topk, hidden_size
    query_entity_embs = tf.reduce_sum(
        topk_ent_embs * tf.expand_dims(topk_softmax, axis=2), axis=1)
    # batch_size, hidden_size
    query_entity_sketch = module.create_cm_sketch(
        entity_ids, topk_softmax, self.all_entity_sketches,
        cm_width=params['cm_width'])
    # batch_size, depth, width
    return query_entity_embs, query_entity_sketch
Compute embeddings and sketch with logits of entities.

    Given entity_ids and logits in the same order, compute weighted average
    of entity embeddings and a sketch that stores their weights.

    Args:
      entity_ids: entity ids
      logits: logits before softmax
      params: params

    Returns:
      A weighted average of embeddings and a cm sketch for the entities
    
language/emql/model.py