def compute_topk_fact_logits(self, topk_fact_idx,
                               query_embs,
                               params):
    """Get the logits between query and facts with topk_fact_idx.

    Args:
      topk_fact_idx: topk ids from scam -- batch_size, topk
      query_embs: embeddings of query -- batch_size, hidden_size * 2
      params: flags

    Returns:
      topk_fact_embs, topk_mask, topk_labels
    """
    topk_subj_ids = tf.gather(self.all_fact_subjids, topk_fact_idx)
    # batch_size, topk
    topk_rel_ids = tf.gather(self.all_fact_relids, topk_fact_idx)
    # batch_size, topk
    topk_mask = tf.cast(topk_subj_ids >= 0, dtype=tf.int32)
    # batch_size, topk

    topk_subj_embs = tf.nn.embedding_lookup(self.entity_embeddings_mat,
                                            topk_subj_ids * topk_mask)
    # batch_size, topk, hidden_size
    topk_rel_embs = tf.nn.embedding_lookup(self.relation_embeddings_mat,
                                           topk_rel_ids * topk_mask)
    # batch_size, topk, hidden_size

    # compute logits for top k facts
    topk_fact_embs = tf.concat((topk_subj_embs, topk_rel_embs), axis=2)
    # batch_size, topk, hidden_size * 2
    topk_fact_embs *= tf.cast(
        tf.expand_dims(topk_mask, axis=2), dtype=tf.float32)
    # batch_size, topk, hidden_size * 2

    query_embs = tf.expand_dims(query_embs, axis=2)
    # batch_size, hidden_size * 2, 1
    topk_fact_logit = tf.matmul(topk_fact_embs, query_embs) / tf.sqrt(
        float(params['entity_emb_size'] + params['relation_emb_size']))
    topk_fact_logit = tf.squeeze(topk_fact_logit, axis=2)
    # batch_size, topk

    return topk_fact_logit
Get the logits between query and facts with topk_fact_idx.

    Args:
      topk_fact_idx: topk ids from scam -- batch_size, topk
      query_embs: embeddings of query -- batch_size, hidden_size * 2
      params: flags

    Returns:
      topk_fact_embs, topk_mask, topk_labels
    
language/emql/model.py