def multi_hop(qry_input_ids,
              qry_input_mask,
              qry_entity_ids,
              entity_ids,
              entity_mask,
              ent2ment_ind,
              ent2ment_val,
              ment2ent_map,
              is_training,
              use_one_hot_embeddings,
              bert_config,
              qa_config,
              mips_config,
              num_hops,
              exclude_set=None,
              bridge_mentions=None,
              answer_mentions=None):
  """Two hops of propagation from input to output entities."""
  # for question BOW embedding
  with tf.variable_scope("qry/bow"):
    word_weights = tf.get_variable(
        "word_weights", [bert_config.vocab_size, 1],
        dtype=tf.float32,
        initializer=tf.ones_initializer())

  qry_seq_emb, word_emb_table = shared_qry_encoder_v2(qry_input_ids,
                                                      qry_input_mask,
                                                      is_training,
                                                      use_one_hot_embeddings,
                                                      bert_config, qa_config)

  batch_size = tf.shape(qry_input_ids)[0]
  if isinstance(qry_entity_ids, tf.SparseTensor):
    # Multiple entities per question. We need to re-score.
    with tf.name_scope("entity_linking"):
      batch_entity_emb = entity_emb(
          tf.cast(qry_entity_ids.values, tf.int64), entity_ids, entity_mask,
          word_emb_table, word_weights)
      qry_el_emb, _ = layer_qry_encoder(
          qry_seq_emb,
          qry_input_ids,
          qry_input_mask,
          is_training,
          bert_config,
          qa_config,
          suffix="_el",
          project=False)
      batch_qry_el_emb = tf.gather(qry_el_emb, qry_entity_ids.indices[:, 0])
      batch_entity_el_scs = tf.reduce_sum(batch_qry_el_emb * batch_entity_emb,
                                          -1)
      batch_entities_nosc = tf.SparseTensor(
          indices=tf.concat([
              qry_entity_ids.indices[:, 0:1],
              tf.cast(tf.expand_dims(qry_entity_ids.values, 1), tf.int64)
          ],
                            axis=1),
          values=batch_entity_el_scs,
          dense_shape=[batch_size, qa_config.num_entities])
      batch_entities = tf.sparse.softmax(tf.sparse.reorder(batch_entities_nosc))
  else:
    batch_entities = tf.SparseTensor(
        indices=tf.concat([
            tf.cast(tf.expand_dims(tf.range(batch_size), 1), tf.int64),
            tf.cast(tf.expand_dims(qry_entity_ids, 1), tf.int64)
        ],
                          axis=1),
        values=tf.ones_like(qry_entity_ids, tf.float32),
        dense_shape=[batch_size, qa_config.num_entities])
  ensure_mentions = bridge_mentions

  with tf.device("/cpu:0"):
    # mips search.
    tf_db, mips_search_fn = search_utils.create_mips_searcher(
        mips_config.ckpt_var_name, mips_config.ckpt_path,
        mips_config.num_neighbors)

  layer_mentions, layer_entities = [], []
  layer_dense, layer_sp = [], []
  for hop in range(num_hops):
    with tf.name_scope("hop_%d" % hop):
      qry_start_emb, qry_end_emb = layer_qry_encoder(
          qry_seq_emb,
          qry_input_ids,
          qry_input_mask,
          is_training,
          bert_config,
          qa_config,
          suffix="_%d" % hop)

      ret_entities, ret_mentions, dense_mention_vec, sp_mention_vec = follow(
          batch_entities, qry_start_emb, qry_end_emb, entity_ids, entity_mask,
          ent2ment_ind, ent2ment_val, ment2ent_map, word_emb_table,
          word_weights, mips_search_fn, tf_db, bert_config.hidden_size,
          mips_config, qa_config, is_training, ensure_mentions)
      if exclude_set is not None:
        batch_ind = tf.expand_dims(tf.range(batch_size), 1)
        exclude_indices = tf.concat([
            tf.cast(batch_ind, tf.int64),
            tf.cast(tf.expand_dims(exclude_set, 1), tf.int64)
        ],
                                    axis=1)
        ret_entities = remove_from_sparse(ret_entities, exclude_indices)
      scaled_entities = tf.sparse.reorder(ret_entities)
      scaled_entities = tf.SparseTensor(
          indices=scaled_entities.indices,
          values=scaled_entities.values / qa_config.softmax_temperature,
          dense_shape=scaled_entities.dense_shape)
      batch_entities = tf.sparse.softmax(scaled_entities)
      ensure_mentions = answer_mentions
      layer_mentions.append(ret_mentions)
      layer_entities.append(ret_entities)
      layer_dense.append(dense_mention_vec)
      layer_sp.append(sp_mention_vec)

  return (layer_entities, layer_mentions, layer_dense, layer_sp,
          batch_entities_nosc, qry_seq_emb)
Two hops of propagation from input to output entities.
language/labs/drkit/model_fns.py