def _knn_graph_blas(x, k, dist="euclidean"):
    r"""Construct a graph from a set of points according to k-nearest-neighbor (KNN).

    This function first compute the distance matrix using BLAS matrix multiplication
    operation provided by backend frameworks. Then use topk algorithm to get
    k-nearest neighbors.

    Parameters
    ----------
    x : Tensor
        The point coordinates. It can be either on CPU or GPU.

        * If is 2D, ``x[i]`` corresponds to the i-th node in the KNN graph.

        * If is 3D, ``x[i]`` corresponds to the i-th KNN graph and
          ``x[i][j]`` corresponds to the j-th node in the i-th KNN graph.
    k : int
        The number of nearest neighbors per node.
    dist : str, optional
        The distance metric used to compute distance between points. It can be the following
        metrics:
        * 'euclidean': Use Euclidean distance (L2 norm) :math:`\sqrt{\sum_{i} (x_{i} - y_{i})^{2}}`.
        * 'cosine': Use cosine distance.
        (default: 'euclidean')
    """
    if F.ndim(x) == 2:
        x = F.unsqueeze(x, 0)
    n_samples, n_points, _ = F.shape(x)

    if k > n_points:
        dgl_warning(
            "'k' should be less than or equal to the number of points in 'x'"
            "expect k <= {0}, got k = {1}, use k = {0}".format(n_points, k)
        )
        k = n_points

    # if use cosine distance, normalize input points first
    # thus we can use euclidean distance to find knn equivalently.
    if dist == "cosine":
        l2_norm = lambda v: F.sqrt(F.sum(v * v, dim=2, keepdims=True))
        x = x / (l2_norm(x) + 1e-5)

    ctx = F.context(x)
    dist = pairwise_squared_distance(x)
    k_indices = F.astype(F.argtopk(dist, k, 2, descending=False), F.int64)
    # index offset for each sample
    offset = F.arange(0, n_samples, ctx=ctx) * n_points
    offset = F.unsqueeze(offset, 1)
    src = F.reshape(k_indices, (n_samples, n_points * k))
    src = F.unsqueeze(src, 0) + offset
    dst = F.repeat(F.arange(0, n_points, ctx=ctx), k, dim=0)
    dst = F.unsqueeze(dst, 0) + offset
    return convert.graph((F.reshape(src, (-1,)), F.reshape(dst, (-1,))))
Construct a graph from a set of points according to k-nearest-neighbor (KNN).

    This function first compute the distance matrix using BLAS matrix multiplication
    operation provided by backend frameworks. Then use topk algorithm to get
    k-nearest neighbors.

    Parameters
    ----------
    x : Tensor
        The point coordinates. It can be either on CPU or GPU.

        * If is 2D, ``x[i]`` corresponds to the i-th node in the KNN graph.

        * If is 3D, ``x[i]`` corresponds to the i-th KNN graph and
          ``x[i][j]`` corresponds to the j-th node in the i-th KNN graph.
    k : int
        The number of nearest neighbors per node.
    dist : str, optional
        The distance metric used to compute distance between points. It can be the following
        metrics:
        * 'euclidean': Use Euclidean distance (L2 norm) :math:`\sqrt{\sum_{i} (x_{i} - y_{i})^{2}}`.
        * 'cosine': Use cosine distance.
        (default: 'euclidean')
    
python/dgl/transforms/functional.py