code: def assign_offsets_pyg(
      node_counts: Dict[str, int], replication_factor: int = 1
  ):
      # cuGraph-PyG assigns offsets based on lexicographic order
      node_offsets = {}
      node_offsets_replicated = {}
      count = 0
      count_replicated = 0
      for node_type in sorted(node_counts.keys()):
          node_offsets[node_type] = count
          node_offsets_replicated[node_type] = count_replicated
  
          count += node_counts[node_type]
          count_replicated += (
              node_counts[node_type] * replication_factor
          )
  
      return node_offsets, node_offsets_replicated, count_replicated
  
imports: import logging
  import warnings
  import argparse
  import traceback
  from cugraph.testing.mg_utils import generate_edgelist_rmat, get_allocation_counts_dask_lazy, sizeof_fmt, get_peak_output_ratio_across_workers, start_dask_client, stop_dask_client, enable_spilling
  from cugraph.structure.symmetrize import symmetrize
  from cugraph.gnn import BulkSampler
  import cugraph
  import json
  import re
  import os
  import gc
  from time import sleep, perf_counter
  from math import ceil
  import pandas
  import numpy
  import cupy
  import cudf
  import dask_cudf
  import dask.dataframe
  from dask.distributed import default_client
  from cugraph.dask import get_n_workers
  from typing import Optional, Union, Dict
label_desc: Use image node_id 7 for calling a global function with example usage: assign_offsets_pyg(node_counts, replication_factor) and returns: node_offsets,  node_offsets_replicated,  count_replicated