def convert_search_to_vector(dist, idx, batch_size, num_nn, vecsize):
  """Create vector from search indices and distances."""
  brange = tf.range(batch_size, dtype=tf.int32)
  bindices = tf.tile(tf.expand_dims(brange, 1), (1, num_nn))
  indices = tf.reshape(
      tf.stack([bindices, idx], axis=2), (batch_size * num_nn, 2))
  values = tf.reshape(dist, (batch_size * num_nn,))
  return tf.SparseTensor(
      indices=tf.cast(indices, dtype=tf.int64),
      values=values,
      dense_shape=[batch_size, vecsize])
Create vector from search indices and distances.
language/labs/drkit/model_fns.py