code: def segmented_knn_graph(
      x,
      k,
      segs,
      algorithm="bruteforce-blas",
      dist="euclidean",
      exclude_self=False,
  ):
      r"""Construct multiple graphs from multiple sets of points according to
      k-nearest-neighbor (KNN) and return.
  
      Compared with :func:`dgl.knn_graph`, this allows multiple point sets with
      different capacity. The points from different sets are stored contiguously
      in the :attr:`x` tensor.
      :attr:`segs` specifies the number of points in each point set. The
      function constructs a KNN graph for each point set, where the predecessors
      of each point are its k-nearest neighbors measured by the Euclidean distance.
      DGL then composes all KNN graphs
      into a batched graph with multiple (:math:`len(segs)`) connected components.
  
      Parameters
      ----------
      x : Tensor
          Coordinates/features of points. Must be 2D. It can be either on CPU or GPU.
      k : int
          The number of nearest neighbors per node.
      segs : list[int]
          Number of points in each point set. The numbers in :attr:`segs`
          must sum up to the number of rows in :attr:`x`.
      algorithm : str, optional
          Algorithm used to compute the k-nearest neighbors.
  
          * 'bruteforce-blas' will first compute the distance matrix
            using BLAS matrix multiplication operation provided by
            backend frameworks. Then use topk algorithm to get
            k-nearest neighbors. This method is fast when the point
            set is small but has :math:`O(N^2)` memory complexity where
            :math:`N` is the number of points.
  
          * 'bruteforce' will compute distances pair by pair and
            directly select the k-nearest neighbors during distance
            computation. This method is slower than 'bruteforce-blas'
            but has less memory overhead (i.e., :math:`O(Nk)` where :math:`N`
            is the number of points, :math:`k` is the number of nearest
            neighbors per node) since we do not need to store all distances.
  
          * 'bruteforce-sharemem' (CUDA only) is similar to 'bruteforce'
            but use shared memory in CUDA devices for buffer. This method is
            faster than 'bruteforce' when the dimension of input points
            is not large. This method is only available on CUDA device.
  
          * 'kd-tree' will use the kd-tree algorithm (CPU only).
            This method is suitable for low-dimensional data (e.g. 3D
            point clouds)
  
          * 'nn-descent' is an approximate approach from paper
            `Efficient k-nearest neighbor graph construction for generic similarity
            measures <https://www.cs.princeton.edu/cass/papers/www11.pdf>`_. This method
            will search for nearest neighbor candidates in "neighbors' neighbors".
  
          (default: 'bruteforce-blas')
      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')
      exclude_self : bool, optional
          If True, the output graph will not contain self loop edges, and each node will not
          be counted as one of its own k neighbors.  If False, the output graph will contain
          self loop edges, and a node will be counted as one of its own k neighbors.
  
      Returns
      -------
      DGLGraph
          The batched graph. The node IDs are in the same order as :attr:`x`.
  
      Examples
      --------
  
      The following examples use PyTorch backend.
  
      >>> import dgl
      >>> import torch
  
      In the example below, the first point set has three points
      and the second point set has four points.
  
      >>> # Features/coordinates of the first point set
      >>> x1 = torch.tensor([[0.0, 0.5, 0.2],
      ...                    [0.1, 0.3, 0.2],
      ...                    [0.4, 0.2, 0.2]])
      >>> # Features/coordinates of the second point set
      >>> x2 = torch.tensor([[0.3, 0.2, 0.1],
      ...                    [0.5, 0.2, 0.3],
      ...                    [0.1, 0.1, 0.2],
      ...                    [0.6, 0.3, 0.3]])
      >>> x = torch.cat([x1, x2], dim=0)
      >>> segs = [x1.shape[0], x2.shape[0]]
      >>> knn_g = dgl.segmented_knn_graph(x, 2, segs)
      >>> knn_g.edges()
      (tensor([0, 0, 1, 1, 1, 2, 3, 3, 4, 4, 5, 5, 6, 6]),
       tensor([0, 1, 0, 1, 2, 2, 3, 5, 4, 6, 3, 5, 4, 6]))
      """
      if exclude_self:
          # add 1 to k, for the self edge, since it will be removed
          k = k + 1
  
      # check invalid k
      if k <= 0:
          raise DGLError(
              "Invalid k value. expect k > 0, got k = {}".format(k)
          )
  
      # check empty point set
      if F.shape(x)[0] == 0:
          raise DGLError("Find empty point set")
  
      if algorithm == "bruteforce-blas":
          result = _segmented_knn_graph_blas(x, k, segs, dist=dist)
      else:
          out = knn(k, x, segs, algorithm=algorithm, dist=dist)
          row, col = out[1], out[0]
          result = convert.graph((row, col))
  
      num_nodes = F.tensor(segs, dtype=F.int64).to(F.context(x))
      result.set_batch_num_nodes(num_nodes)
      # if any segment is too small for k, all algorithms reduce k for all segments
      clamped_k = min(k, np.min(segs))
      result.set_batch_num_edges(clamped_k * num_nodes)
  
      if exclude_self:
          # remove_self_loop will update batch_num_edges as needed
          result = remove_self_loop(result)
  
          # If there were more than k(+1) coincident points, there may not have been self loops on
          # all nodes, in which case there would still be one too many out edges on some nodes.
          # However, if every node had a self edge, the common case, every node would still have the
          # same degree as each other, so we can check that condition easily.
          # The -1 is for the self edge removal.
          clamped_k = min(k, np.min(segs)) - 1
          if result.num_edges() != clamped_k * result.num_nodes():
              # edges on any nodes with too high degree should all be length zero,
              # so pick an arbitrary one to remove from each such node
              degrees = result.in_degrees()
              node_indices = F.nonzero_1d(degrees > clamped_k)
              edges_to_remove_graph = sample_neighbors(
                  result, node_indices, 1, edge_dir="in"
              )
              edge_ids = edges_to_remove_graph.edata[EID]
              result = remove_edges(result, edge_ids)
  
      return result
  
imports: import copy
  from collections.abc import Iterable, Mapping
  import numpy
  import scipy.sparse
  import scipy.sparse.linalg
  from ..utils import version
  from ..None import backend, batch, convert, function, ndarray, subgraph, utils
  from .._ffi.function import _init_api
  from ..base import dgl_warning, DGLError, EID, NID
  from ..frame import Frame
  from ..heterograph import DGLGraph
  from ..heterograph_index import create_heterograph_from_relations, create_metagraph_index
  from ..partition import metis_partition, metis_partition_assignment, partition_graph_with_halo
  from ..sampling.neighbor import sample_neighbors
label_desc: Use image node_id 4 for calling a global function with example usage: segmented_knn_graph(x, k, segs, algorithm, dist, exclude_) and returns: result