code: def get_long_tail_partition(n_relations, n_machine):
      """Relation types has a long tail distribution for many dataset.
      So we need to average shuffle the data before we partition it.
      """
      assert n_relations > 0, "n_relations must be a positive number."
      assert n_machine > 0, "n_machine must be a positive number."
  
      partition_book = [0] * n_relations
  
      part_id = 0
      for i in range(n_relations):
          partition_book[i] = part_id
          part_id += 1
          if part_id == n_machine:
              part_id = 0
  
      return partition_book
  
imports: import os
  import argparse
  import time
  import logging
  import socket
  import dgl
  import dgl.backend
  import torch.multiprocessing
  from .train_pytorch import load_model, dist_train_test
  from .utils import get_compatible_batch_size, CommonArgParser
  from .train import prepare_save_path
  from .dataloader import TrainDataset, NewBidirectionalOneShotIterator
  from .dataloader import get_partition_dataset
label_desc: Use image node_id 1 for calling a global function with example usage: get_long_tail_partition(n_relations, n_machine) and returns: partition_book