def model_metaqa(self, features, params,
                   hop, top_k):
    """Compute logits for MetaQA multi-hop reasoning task.

    MetaQA model is made of 4 different parts:
    1. Construct KB triple query to retrieve the top k triples from KB.
    2. Check if retrieved triples are eligible by filtering them with sketch.
    3. Check if any retrieved triples should be excluded, i.e. if their
       object entities have been visited in the previous iterations.
    4. Update sketch and query embeddings from the next iteration.

    Args:
      features:
        -- question: a list of token ids for vocabs padded with -1
        -- question_entity_id: question entity id
        -- question_entity_sketch: a count-min sketch with only one element
                                   (the quesion entity)
        -- answer_labels: a k-hot vector marking if an entity is an answer
      params: hyper-parameters
      hop: number of hops
      top_k: select top_k facts at each iteration

    Returns:
      predictions and losses
    """
    tensors = {}
    question, question_entity_id, question_entity_sketch, answer_labels = features

    # Get text encoding of questions and question question entities.
    question_embs = self.compute_average_embs(question)
    # batch_size, hidden_size
    question_entity_embs = tf.nn.embedding_lookup(
        self.entity_embeddings_mat, question_entity_id)
    # batch_size, hidden_size

    # Set initial query embeddings and sketchs.
    query_entity_embs = question_entity_embs
    # batch_size, hidden_size
    query_entity_sketch = question_entity_sketch
    # batch_size, depth, width
    excluded_entity_sketch = question_entity_sketch
    # batch_size, depth, width

    # Begin iteration.
    for i in range(hop):
      # Construct queries by mapping original question embeddings to another
      # space. Queries are constructed as a concatenation of subject, relation,
      # and placeholders for object embeddings.
      with tf.variable_scope('question_emb_ffn_%s' % i):
        question_embs = tf.keras.layers.Dense(
            units=params['relation_emb_size'])(
                question_embs)
        # batch_size, hidden_size
        query_embs = tf.concat(
            [query_entity_embs, question_embs,
             tf.zeros_like(query_entity_embs)],
            axis=1)
        # batch_size, hidden_size * 3

      # Retrieve top k facts that match the queries, and check if the top
      # k facts are eligible by checking the count-min sketch. Subject
      # entities must have non-zero probability to be considered as
      # eligible.
      topk_fact_logits, topk_fact_ids = util.topk_search_fn(
          query_embs, self.kb_embeddings_mat, top_k)
      # batch_size, top_k     batch_size, top_k
      topk_subj_ids = tf.gather(self.all_fact_subjids, topk_fact_ids)
      # batch_size, topk
      is_topk_fact_eligible, _ = module.check_topk_fact_eligible(
          topk_subj_ids, query_entity_sketch, self.all_entity_sketches, params)

      # Entities visited before should also be excluded. Similarly, we first
      # get the probability from sketch and exclude those with non-zero
      # probability in previous iterations.
      topk_obj_ids = tf.gather(self.all_fact_objids, topk_fact_ids)
      # batch_size, topk
      is_topk_fact_excluded, _ = module.check_topk_fact_eligible(
          topk_obj_ids, excluded_entity_sketch,
          self.all_entity_sketches, params)

      # Apply the filtering results to logits of topk facts.
      if params['use_cm_sketch']:
        topk_fact_logits += (
            1 - is_topk_fact_eligible + is_topk_fact_excluded) * VERY_NEG
        # batch_size, top_k

      # Construct query embeddings for next iteration.
      topk_fact_obj_embs = tf.nn.embedding_lookup(
          self.entity_embeddings_mat, topk_obj_ids)
      # batch_size, top_k, hidden_size
      topk_softmax = tf.nn.softmax(topk_fact_logits)
      # batch_size, top_k, 1
      query_entity_embs = tf.reduce_sum(
          topk_fact_obj_embs * tf.expand_dims(topk_softmax, axis=2), axis=1)
      # batch_size, hidden_size

      # Update sketches.
      query_entity_sketch = module.create_cm_sketch(
          topk_obj_ids, topk_softmax, self.all_entity_sketches,
          cm_width=params['cm_width'])
      # batch_size, depth, width
      excluded_entity_sketch += query_entity_sketch
      # batch_size, depth, width

    # We only compute loss on the topk retrieval results at the last iteration.
    # No intermediate training signal is required.
    topk_fact_labels = tf.gather(
        answer_labels, topk_fact_ids, batch_dims=1, axis=1)
    topk_fact_loss = self.compute_tf_loss(topk_fact_logits, topk_fact_labels)
    tensors = {
        'logits': topk_fact_logits,
        'labels': topk_fact_labels,
        'candidates': topk_fact_ids
    }
    return topk_fact_loss, tensors
Compute logits for MetaQA multi-hop reasoning task.

    MetaQA model is made of 4 different parts:
    1. Construct KB triple query to retrieve the top k triples from KB.
    2. Check if retrieved triples are eligible by filtering them with sketch.
    3. Check if any retrieved triples should be excluded, i.e. if their
       object entities have been visited in the previous iterations.
    4. Update sketch and query embeddings from the next iteration.

    Args:
      features:
        -- question: a list of token ids for vocabs padded with -1
        -- question_entity_id: question entity id
        -- question_entity_sketch: a count-min sketch with only one element
                                   (the quesion entity)
        -- answer_labels: a k-hot vector marking if an entity is an answer
      params: hyper-parameters
      hop: number of hops
      top_k: select top_k facts at each iteration

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