code: def data_partition(dataset_x, dataset_y, global_rank, world_size):
      data_per_rank = dataset_x.shape[0] // world_size
      idx_start = global_rank * data_per_rank
      idx_end = (global_rank + 1) * data_per_rank
      return dataset_x[idx_start:idx_end], dataset_y[idx_start:idx_end]
  
imports: from singa import singa_wrap
  from singa import autograd
  from singa import tensor
  from singa import device
  from singa import opt
  from PIL import Image
  import numpy
  import os
  import sys
  import time
label_desc: Use image node_id 12 for calling a global function with example usage: data_partition(dataset_x, dataset_y, global_rank, world_size) and returns: dataset_x,  dataset_y