def create_mips_searcher(var_name,
                         checkpoint_path,
                         num_neighbors,
                         local_var_name="mips_init_barrier"):
  """Create searcher for returning top-k closest elements."""
  tf_db = load_database(var_name, None, checkpoint_path)

  with tf.control_dependencies([tf_db.initializer]):
    mips_init_barrier = tf.constant(True)

  # Make sure DB is initialized.
  tf.get_local_variable(local_var_name, initializer=mips_init_barrier)

  def _search(query):
    with tf.device("/cpu:0"):
      distance = tf.matmul(query, tf_db, transpose_b=True)
      topk_dist, topk_idx = tf.nn.top_k(distance, num_neighbors)
    topk_dist.set_shape([query.shape[0], num_neighbors])
    topk_idx.set_shape([query.shape[0], num_neighbors])
    return topk_dist, topk_idx

  return tf_db, _search
Create searcher for returning top-k closest elements.
language/labs/drkit/search_utils.py