def rescore_sparse(sp_mentions, tf_db, qrys):
  """Rescore mentions in sparse tensor with given batch of queries."""
  batch_ind = sp_mentions.indices[:, 0]
  mention_ind = sp_mentions.indices[:, 1]
  mention_vec = tf.gather(tf_db, mention_ind)
  batch_qrys = tf.gather(qrys, batch_ind)
  mention_scs = tf.reduce_sum(batch_qrys * mention_vec, 1)
  return tf.SparseTensor(
      indices=sp_mentions.indices,
      values=mention_scs,
      dense_shape=sp_mentions.dense_shape)
Rescore mentions in sparse tensor with given batch of queries.
language/labs/drkit/model_fns.py