code: def infer_input_nodes(
      data: Union[Data, HeteroData], split: str
  ) -> InputNodes:
      attr_name: Optional[str] = None
      if f"{split}_mask" in data:
          attr_name = f"{split}_mask"
      elif f"{split}_idx" in data:
          attr_name = f"{split}_idx"
      elif f"{split}_index" in data:
          attr_name = f"{split}_index"
  
      if attr_name is None:
          return None
  
      if isinstance(data, Data):
          return data[attr_name]
      if isinstance(data, HeteroData):
          input_nodes_dict = {
              node_type: store[attr_name]
              for node_type, store in data.node_items()
              if attr_name in store
          }
          if len(input_nodes_dict) != 1:
              raise ValueError(
                  f"Could not automatically determine the input "
                  f"nodes of {data} since there exists multiple "
                  f"types with attribute '{attr_name}'"
              )
          return list(input_nodes_dict.items())[0]
      return None
  
imports: import copy
  import inspect
  import warnings
  from typing import Any, Dict, Optional, Tuple, Type, Union
  import torch
  from torch_geometric.data import Data, Dataset, HeteroData
  from torch_geometric.loader import DataLoader, LinkLoader, NodeLoader
  from torch_geometric.sampler import BaseSampler, NeighborSampler
  from torch_geometric.typing import InputEdges, InputNodes, OptTensor
label_desc: Use image node_id 1 for calling a global function with example usage: infer_input_nodes(data, split) and returns: None, None, data, list