code: def create_mask_per_rank(
      global_mask: Union[torch.Tensor, Tuple[str, torch.Tensor]],
      rank: int,
      world_size: int,
      hetero: bool = False,
  ) -> Union[torch.Tensor, Tuple[str, torch.Tensor]]:
      mask = global_mask[-1] if hetero else global_mask
      nonzero = mask.nonzero().reshape(-1)
      rank_indices = nonzero.split(
          nonzero.size(0) // world_size, dim=0
      )[rank].clone()
      mask_per_rank = torch.full_like(mask, False)
      mask_per_rank[rank_indices] = True
  
      if hetero:
          return tuple((global_mask[0], mask_per_rank))
      else:
          return mask_per_rank
  
file_path: pytorch_geometric/benchmark/multi_gpu/training/common.py
imports: import argparse
  import ast
  from time import perf_counter
  from typing import Any, Callable, Tuple, Union
  import torch
  import torch.distributed
  import torch.nn.functional
  from torch.nn.parallel import DistributedDataParallel
  from benchmark.utils import get_model, get_split_masks, test
  from torch_geometric.data import Data, HeteroData
  from torch_geometric.loader import NeighborLoader
  from torch_geometric.nn import PNAConv
label_desc: Use image node_id 4 for calling a global function with example usage: create_mask_per_rank(global_mask, rank, world_size, hetero) and returns: tuple, mask_per_rank