class_docstr: The graph attentional operator on heterogeneous graphs, where a separate
  `GATConv` is applied on the homogeneous graph for each edge type. Compared
  with directly wrapping `GATConv`s with `HeteroConv`, `HeteroGATConv` fuses
  all the linear transformation associated with each node type together into 1
  GEMM call, to improve the performance on GPUs.
  
  Parameters
  ----------
  in_channels : int or Dict[str, int])
      Size of each input sample of every node type.
  
  out_channels : int
      Size of each output sample.
  
  node_types : List[str]
      List of Node types.
  
  edge_types : List[Tuple[str, str, str]]
      List of Edge types.
  
  heads : int, optional (default=1)
      Number of multi-head-attentions.
  
  concat : bool, optional (default=True):
      If set to :obj:`False`, the multi-head attentions are averaged instead
      of concatenated.
  
  negative_slope : float, optional (default=0.2)
      LeakyReLU angle of the negative slope.
  
  bias : bool, optional (default=True)
      If set to :obj:`False`, the layer will not learn an additive bias.
  
  aggr : str, optional (default="sum")
      The aggregation scheme to use for grouping node embeddings generated by
      different relations. Choose from "sum", "mean", "min", "max".
class_docstr_tok: ["The","graph","attentional","operator","on","heterogeneous","graphs",",","where","a","separate","`","GATConv","`","is","applied","on","the","homogeneous","graph","for","each","edge","type",".","Compared","with","directly","wrapping","`","GATConv","`","s","with","`","HeteroConv","`",",","`","HeteroGATConv","`","fuses","all","the","linear","transformation","associated","with","each","node","type","together","into","1","GEMM","call",",","to","improve","the","performance","on","GPUs",".","Parameters","--","--","--","--","--","in_channels",":","int","or","Dict","[","str",",","int","]",")","Size","of","each","input","sample","of","every","node","type",".","out_channels",":","int","Size","of","each","output","sample",".","node_types",":","List","[","str","]","List","of","Node","types",".","edge_types",":","List","[","Tuple","[","str",",","str",",","str","]","]","List","of","Edge","types",".","heads",":","int",",","optional","(","default=1",")","Number","of","multi-head-attentions",".","concat",":","bool",",","optional","(","default=True",")",":","If","set","to",":","obj",":","`","False","`",",","the","multi-head","attentions","are","averaged","instead","of","concatenated",".","negative_slope",":","float",",","optional","(","default=0.2",")","LeakyReLU","angle","of","the","negative","slope",".","bias",":","bool",",","optional","(","default=True",")","If","set","to",":","obj",":","`","False","`",",","the","layer","will","not","learn","an","additive","bias",".","aggr",":","str",",","optional","(","default=","''","sum","''",")","The","aggregation","scheme","to","use","for","grouping","node","embeddings","generated","by","different","relations",".","Choose","from","``","sum","''",",","``","mean","''",",","``","min","''",",","``","max","''","."]
docstr: Split fused tensors into chunks based on edge types.
  
  Parameters
  ----------
  x_fused_dict : dict[str, torch.Tensor]
      A dictionary to hold node feature for each node type. The key is
      node type; the value is a fused tensor that account for all
      relations for that node type.
  
  dim : int
      Dimension along which to split the fused tensor.
  
  Returns
  -------
  x_src_dict : dict[str, torch.Tensor]
      A dictionary to hold source node feature for each relation graph.
  
  x_dst_dict : dict[str, torch.Tensor]
      A dictionary to hold destination node feature for each relation graph.
docstr_tok: ["Split","fused","tensors","into","chunks","based","on","edge","types",".","Parameters","--","--","--","--","--","x_fused_dict",":","dict","[","str",",","torch.Tensor","]","A","dictionary","to","hold","node","feature","for","each","node","type",".","The","key","is","node","type",";","the","value","is","a","fused","tensor","that","account","for","all","relations","for","that","node","type",".","dim",":","int","Dimension","along","which","to","split","the","fused","tensor",".","Returns","--","--","--","-","x_src_dict",":","dict","[","str",",","torch.Tensor","]","A","dictionary","to","hold","source","node","feature","for","each","relation","graph",".","x_dst_dict",":","dict","[","str",",","torch.Tensor","]","A","dictionary","to","hold","destination","node","feature","for","each","relation","graph","."]
code: def split_tensors(
      self, x_fused_dict: dict[str, torch.Tensor], dim: int
  ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
      """Split fused tensors into chunks based on edge types.
  
      Parameters
      ----------
      x_fused_dict : dict[str, torch.Tensor]
          A dictionary to hold node feature for each node type. The key is
          node type; the value is a fused tensor that account for all
          relations for that node type.
  
      dim : int
          Dimension along which to split the fused tensor.
  
      Returns
      -------
      x_src_dict : dict[str, torch.Tensor]
          A dictionary to hold source node feature for each relation graph.
  
      x_dst_dict : dict[str, torch.Tensor]
          A dictionary to hold destination node feature for each relation graph.
      """
      x_src_dict = dict.fromkeys(self.edge_types)
      x_dst_dict = dict.fromkeys(self.edge_types)
  
      for ntype, t in x_fused_dict.items():
          n_src_rel = len(self.relations_per_ntype[ntype][0])
          n_dst_rel = len(self.relations_per_ntype[ntype][1])
          n_rel = n_src_rel + n_dst_rel
          t_list = torch.chunk(t, chunks=n_rel, dim=dim)
  
          for i, src_rel in enumerate(
              self.relations_per_ntype[ntype][0]
          ):
              x_src_dict[src_rel] = t_list[i]
  
          for i, dst_rel in enumerate(
              self.relations_per_ntype[ntype][1]
          ):
              x_dst_dict[dst_rel] = t_list[i + n_src_rel]
  
      return x_src_dict, x_dst_dict
  
file_path: cugraph/python/cugraph-pyg/cugraph_pyg/nn/conv/hetero_gat_conv.py
imports: from typing import Optional, Union
  from collections import defaultdict
  from cugraph.utilities.utils import import_optional
  from pylibcugraphops.pytorch.operators import mha_gat_n2n
  from .base import BaseConv
label_desc: Use image node_id 2 for calling the HeteroGATConv obj's underlying member method code with example usage: obj.split_tensors(x_fused_dict, dim) and returns: x_src_dict,  x_dst_dict