tensorflow.py 文件源码

def __knn_bruteforce(X, k=50):
    # calculate euclidean distances
    r = tf.reduce_sum(X*X, 1)
    r = tf.reshape(r, [-1, 1])
    D = r - 2*tf.matmul(X, tf.transpose(X)) + tf.transpose(r)
    D = tf.matrix_set_diag(D, tf.constant(1e32, dtype=X.dtype,
                           shape=(X.shape[0],)))
    D = tf.sqrt(D)

    #find kNNs
    distances, indices = tf.nn.top_k(-D, k)
    return -distances, indices