code: def group_hetero_graph(
      edge_index_dict: Dict[EdgeType, Tensor],
      num_nodes_dict: Optional[Dict[NodeType, int]] = None,
  ) -> Tuple[
      Tensor,
      Tensor,
      Tensor,
      Tensor,
      Dict[Union[str, int], Tensor],
      Dict[Union[NodeType, EdgeType], int],
  ]:
      num_nodes_dict = maybe_num_nodes_dict(
          edge_index_dict, num_nodes_dict
      )
  
      tmp = list(edge_index_dict.values())[0]
  
      key2int: Dict[Union[NodeType, EdgeType], int] = {}
  
      cumsum, offset = 0, {}  # Helper data.
      node_types, local_node_indices = [], []
      local2global: Dict[Union[str, int], Tensor] = {}
      for i, (key, N) in enumerate(num_nodes_dict.items()):
          key2int[key] = i
          node_types.append(tmp.new_full((N,), i))
          local_node_indices.append(torch.arange(N, device=tmp.device))
          offset[key] = cumsum
          local2global[key] = local_node_indices[-1] + cumsum
          local2global[i] = local2global[key]
          cumsum += N
  
      node_type = torch.cat(node_types, dim=0)
      local_node_idx = torch.cat(local_node_indices, dim=0)
  
      edge_indices, edge_types = [], []
      for i, (keys, edge_index) in enumerate(edge_index_dict.items()):
          key2int[keys] = i
          inc = torch.tensor([offset[keys[0]], offset[keys[-1]]]).view(
              2, 1
          )
          edge_indices.append(edge_index + inc.to(tmp.device))
          edge_types.append(tmp.new_full((edge_index.size(1),), i))
  
      edge_index = torch.cat(edge_indices, dim=-1)
      edge_type = torch.cat(edge_types, dim=0)
  
      return (
          edge_index,
          edge_type,
          node_type,
          local_node_idx,
          local2global,
          key2int,
      )
  
file_path: pytorch_geometric/torch_geometric/utils/hetero.py
imports: from typing import Dict, List, Optional, Set, Tuple, Union
  import torch
  from torch import Tensor
  from torch.nn import ParameterDict
  from torch_geometric.typing import Adj, EdgeType, NodeType, SparseTensor
  from torch_geometric.utils import is_sparse, to_edge_index
  from torch_geometric.utils.num_nodes import maybe_num_nodes_dict
label_desc: Use image node_id 1 for calling a global function with example usage: group_hetero_graph(edge_index_dict, num_nodes_dict) and returns: edge_index,  edge_type,  node_type,  local_node_idx,  local2global,  key2int