def model_follow(self, features,
                   params):
    """Compute logits for follow operation.

    A follow operation is considered as the intersection of a set of
    facts with correct subjects and a set of facts with correct relation.

    Args:
      features: features
        -- subject_set: all facts that have the correct subject
        -- relation_set: all facts that have the correct relation
        -- labels: facts that have both correct subject and relation
      params: hyper-parameters

    Returns:
      predictions and losses
    """
    subject_set, relation_set, labels = features
    subject_set_mask = tf.expand_dims(subject_set, axis=2)
    # batch_size, num_candidate, 1
    relation_set_mask = tf.expand_dims(relation_set, axis=2)
    # batch_size, num_candidate, 1
    candidate_embs = self.load_kb_embeddings('set_follow', params)
    # batch_size, num_candidate, hidden_size * 3

    set1_embs_query = tf.reduce_sum(
        candidate_embs * subject_set_mask, axis=1, keepdims=True) \
        / tf.reduce_sum(subject_set_mask, axis=1, keepdims=True)
    # batch_size, 1, hidden_size * 3
    set2_embs_query = tf.reduce_sum(
        candidate_embs * relation_set_mask, axis=1, keepdims=True) \
        / tf.reduce_sum(relation_set_mask, axis=1, keepdims=True)
    # batch_size, 1, hidden_size * 3
    intersection_embs_query = set1_embs_query + set2_embs_query
    # batch_size, 1, hidden_size * 3
    logits, _ = self.compute_logits(
        candidate_embs, intersection_embs_query,
        params['entity_emb_size'] * 2 + params['relation_emb_size'])
    loss = self.compute_tf_loss(logits, labels)
    return loss, {'logits': logits, 'labels': labels}
Compute logits for follow operation.

    A follow operation is considered as the intersection of a set of
    facts with correct subjects and a set of facts with correct relation.

    Args:
      features: features
        -- subject_set: all facts that have the correct subject
        -- relation_set: all facts that have the correct relation
        -- labels: facts that have both correct subject and relation
      params: hyper-parameters

    Returns:
      predictions and losses
    
language/emql/model.py