code: def get_mask_size(
      attr: str, store: BaseStorage, size: Optional[int]
  ) -> int:
      if size is not None:
          return size
      return (
          store.num_edges
          if store.is_edge_attr(attr)
          else store.num_nodes
      )
  
file_path: pytorch_geometric/torch_geometric/transforms/mask.py
imports: from typing import List, Optional, Sequence, Union
  from torch_geometric.data import Data, HeteroData
  from torch_geometric.data.datapipes import functional_transform
  from torch_geometric.data.storage import BaseStorage
  from torch_geometric.transforms import BaseTransform
  from torch_geometric.utils import index_to_mask, mask_to_index
label_desc: Use image node_id 2 for calling a global function with example usage: get_mask_size(attr, store, size) and returns: unknown, size