code: def fused_forward(
      self,
      x: Tensor,
      edge_index: Tensor,
      node_type: Tensor,
      edge_type: Tensor,
  ) -> Tensor:
      # TODO This currently does not fuse at all :(
      # TODO We currently assume that `x` and `edge_index` are both sorted
      # according to `type`.
  
      node_sizes = scatter(
          torch.ones_like(node_type),
          node_type,
          dim=0,
          dim_size=len(self.node_types),
          reduce="sum",
      )
      edge_sizes = scatter(
          torch.ones_like(edge_type),
          edge_type,
          dim=0,
          dim_size=len(self.edge_types),
          reduce="sum",
      )
  
      ptr = cumsum(node_sizes)
  
      xs = x.split(node_sizes.tolist())
      x_dict = {
          node_type: x for node_type, x in zip(self.node_types, xs)
      }
  
      # TODO Consider out-sourcing to its own function.
      edge_indices = edge_index.clone().split(
          edge_sizes.tolist(), dim=1
      )
      for (src, _, dst), index in zip(self.edge_types, edge_indices):
          index[0] -= ptr[self.node_type_to_index[src]]
          index[1] -= ptr[self.node_type_to_index[dst]]
  
      edge_index_dict = {
          edge_type: edge_index
          for edge_type, edge_index in zip(
              self.edge_types, edge_indices
          )
      }
  
      out_dict = self.hetero_module(x_dict, edge_index_dict)
      return torch.cat(
          [out_dict[key] for key in self.node_types], dim=0
      )
  
imports: import copy
  import warnings
  from typing import Dict, List, Optional, Union
  import torch
  import torch.nn.functional
  from torch import Tensor
  import torch_geometric
  from torch_geometric import is_compiling
  from torch_geometric.typing import EdgeType, NodeType, OptTensor
  from torch_geometric.utils import cumsum, scatter
label_desc: Use image node_id 2 for calling the ToHeteroMessagePassing obj's underlying member method code with example usage: obj.fused_forward(x, edge_index, node_type, edge_type) and returns: torch