code: def unique_and_compact_csc_formats(
      csc_formats: Union[
          Tuple[torch.Tensor, torch.Tensor],
          Dict[str, Tuple[torch.Tensor, torch.Tensor]],
      ],
      unique_dst_nodes: Union[
          torch.Tensor,
          Dict[str, torch.Tensor],
      ],
  ):
      """
      Compact csc formats and return unique nodes (per type).
  
      Parameters
      ----------
      csc_formats : Union[CSCFormatBase, Dict(str, CSCFormatBase)]
          CSC formats representing source-destination edges.
          - If `csc_formats` is a CSCFormatBase: It means the graph is
          homogeneous. Also, indptr and indice in it should be torch.tensor
          representing source and destination pairs in csc format. And IDs inside
          are homogeneous ids.
          - If `csc_formats` is a Dict[str, CSCFormatBase]: The keys
          should be edge type and the values should be csc format node pairs.
          And IDs inside are heterogeneous ids.
      unique_dst_nodes: torch.Tensor or Dict[str, torch.Tensor]
          Unique nodes of all destination nodes in the node pairs.
          - If `unique_dst_nodes` is a tensor: It means the graph is homogeneous.
          - If `csc_formats` is a dictionary: The keys are node type and the
          values are corresponding nodes. And IDs inside are heterogeneous ids.
  
      Returns
      -------
      Tuple[csc_formats, unique_nodes]
          The compacted csc formats, where node IDs are replaced with mapped node
          IDs, and the unique nodes (per type).
          "Compacted csc formats" indicates that the node IDs in the input node
          pairs are replaced with mapped node IDs, where each type of node is
          mapped to a contiguous space of IDs ranging from 0 to N.
  
      Examples
      --------
      >>> import dgl.graphbolt as gb
      >>> N1 = torch.LongTensor([1, 2, 2])
      >>> N2 = torch.LongTensor([5, 5, 6])
      >>> unique_dst = {
      ...     "n1": torch.LongTensor([1, 2]),
      ...     "n2": torch.LongTensor([5, 6])}
      >>> csc_formats = {
      ...     "n1:e1:n2": gb.CSCFormatBase(indptr=torch.tensor([0, 2, 3]),indices=N1),
      ...     "n2:e2:n1": gb.CSCFormatBase(indptr=torch.tensor([0, 1, 3]),indices=N2)}
      >>> unique_nodes, compacted_csc_formats = gb.unique_and_compact_csc_formats(
      ...     csc_formats, unique_dst
      ... )
      >>> print(unique_nodes)
      {'n1': tensor([1, 2]), 'n2': tensor([5, 6])}
      >>> print(compacted_csc_formats)
      {"n1:e1:n2": CSCFormatBase(indptr=torch.tensor([0, 2, 3]),
                                 indices=torch.tensor([0, 1, 1])),
       "n2:e2:n1": CSCFormatBase(indptr=torch.tensor([0, 1, 3]),
                                 indices=torch.Longtensor([0, 0, 1]))}
      """
      is_homogeneous = not isinstance(csc_formats, dict)
      if is_homogeneous:
          csc_formats = {"_N:_E:_N": csc_formats}
          if unique_dst_nodes is not None:
              assert isinstance(
                  unique_dst_nodes, torch.Tensor
              ), "Edge type not supported in homogeneous graph."
              unique_dst_nodes = {"_N": unique_dst_nodes}
  
      # Collect all source and destination nodes for each node type.
      indices = defaultdict(list)
      device = None
      for etype, csc_format in csc_formats.items():
          if device is None:
              device = csc_format.indices.device
          src_type, _, dst_type = etype_str_to_tuple(etype)
          assert len(unique_dst_nodes.get(dst_type, [])) + 1 == len(
              csc_format.indptr
          ), "The seed nodes should correspond to indptr."
          indices[src_type].append(csc_format.indices)
      indices = {
          ntype: torch.cat(nodes) for ntype, nodes in indices.items()
      }
  
      ntypes = set(indices.keys())
      unique_nodes = {}
      compacted_indices = {}
      dtype = list(indices.values())[0].dtype
      default_tensor = torch.tensor([], dtype=dtype, device=device)
      for ntype in ntypes:
          indice = indices.get(ntype, default_tensor)
          unique_dst = unique_dst_nodes.get(ntype, default_tensor)
          (
              unique_nodes[ntype],
              compacted_indices[ntype],
              _,
          ) = torch.ops.graphbolt.unique_and_compact(
              indice,
              torch.tensor([], dtype=indice.dtype, device=device),
              unique_dst,
          )
  
      compacted_csc_formats = {}
      # Map back with the same order.
      for etype, csc_format in csc_formats.items():
          num_elem = csc_format.indices.size(0)
          src_type, _, _ = etype_str_to_tuple(etype)
          indice = compacted_indices[src_type][:num_elem]
          indptr = csc_format.indptr
          compacted_csc_formats[etype] = CSCFormatBase(
              indptr=indptr, indices=indice
          )
          compacted_indices[src_type] = compacted_indices[src_type][
              num_elem:
          ]
  
      # Return singleton for a homogeneous graph.
      if is_homogeneous:
          compacted_csc_formats = list(compacted_csc_formats.values())[
              0
          ]
          unique_nodes = list(unique_nodes.values())[0]
  
      return unique_nodes, compacted_csc_formats
  
file_path: dgl/python/dgl/graphbolt/internal/sample_utils.py
imports: import copy
  from collections import defaultdict
  from typing import Dict, List, Optional, Tuple, Union
  import torch
  from ..base import CSCFormatBase, etype_str_to_tuple
label_desc: Use image node_id 3 for calling a global function with example usage: unique_and_compact_csc_formats(csc_formats, unique_dst_nodes) and returns: unique_nodes,  compacted_csc_formats