def model_membership(self, features,
                       unused_params):
    """Compute logits for set membership.

    A representation of set is the average of embeddings of entities in
    the set.

    Args:
      features:
        -- entity_ids: ids of entities in the set, padded with -1
                       e.g. [1, 5, 3, -1, -1]
                       # batch_size, max_set
        -- labels: if the i'th entity is in the set, 1 for True, 0 for False
                   e.g. [0, 1, 0, 1, 0, 1, 0, 0, ...]
                   # batch_size, num_entities
      unused_params: hyper-parameters

    Returns:
      predictions and losses
    """
    entity_ids, labels = features
    entity_mask = tf.cast(tf.not_equal(entity_ids, -1), dtype=tf.int32)
    # batch_size, max_set
    entity_embs = tf.nn.embedding_lookup(
        self.entity_embeddings_mat, entity_ids * entity_mask)
    # batch_size, max_set, hidden_size
    entity_mask_3d = tf.cast(
        tf.expand_dims(entity_mask, axis=2), dtype=tf.float32)
    # batch_size, max_set, 1
    set_embs = tf.reduce_sum(entity_embs * entity_mask_3d, axis=1)
    # batch_size, hidden_size
    set_embs /= tf.reduce_sum(entity_mask_3d, axis=1)
    # batch_size, hidden_size
    logits = tf.matmul(set_embs, self.entity_embeddings_mat, transpose_b=True)
    # batch_size, num_entities
    loss = self.compute_tf_loss(logits, labels)
    return loss, {'logits': logits, 'labels': labels}
Compute logits for set membership.

    A representation of set is the average of embeddings of entities in
    the set.

    Args:
      features:
        -- entity_ids: ids of entities in the set, padded with -1
                       e.g. [1, 5, 3, -1, -1]
                       # batch_size, max_set
        -- labels: if the i'th entity is in the set, 1 for True, 0 for False
                   e.g. [0, 1, 0, 1, 0, 1, 0, 0, ...]
                   # batch_size, num_entities
      unused_params: hyper-parameters

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