code: def partition(
      global_rank, world_size, train_x, train_y, val_x, val_y
  ):
      # Partition training data
      data_per_rank = train_x.shape[0] // world_size
      idx_start = global_rank * data_per_rank
      idx_end = (global_rank + 1) * data_per_rank
      train_x = train_x[idx_start:idx_end]
      train_y = train_y[idx_start:idx_end]
  
      # Partition evaluation data
      data_per_rank = val_x.shape[0] // world_size
      idx_start = global_rank * data_per_rank
      idx_end = (global_rank + 1) * data_per_rank
      val_x = val_x[idx_start:idx_end]
      val_y = val_y[idx_start:idx_end]
      return train_x, train_y, val_x, val_y
  
file_path: singa/examples/cnn/train_cnn.py
imports: from singa import singa_wrap
  from singa import device
  from singa import tensor
  from singa import opt
  import numpy
  import time
  import argparse
  from PIL import Image
label_desc: Use image node_id 3 for calling a global function with example usage: partition(global_rank, world_size, train_x, train_y, val_x, val_y) and returns: train_x,  train_y,  val_x,  val_y