class_docstr: Create and run the model. We try with both storing indices for gather
  on CPU and on GPU
code: def run_model(self, V, gpu_devices):
      def input_builder_fun(model):
          return None
  
      def model_build_fun(model, loss_scale):
          gpu_vecs_gathered = []
          gpu_vecs = []
          for num, vec in enumerate(self.vecs):
              gpu_vec = model.param_init_net.CopyCPUToGPU(
                  vec,
                  "gpuvec_{}".format(num),
              )
              if num != 2:
                  model.params.append(gpu_vec)
              gpu_vecs.append(gpu_vec)
          for num, gpu_vec in enumerate(gpu_vecs):
              gpu_vec_gathered = model.net.Gather(
                  [gpu_vec, "indices"],
                  ["gpu_vec_gathered_{}".format(num)],
              )
              gpu_vecs_gathered.append(gpu_vec_gathered)
  
          assert len(gpu_vecs_gathered) == 3
  
          fc = model.net.FC(
              [
                  gpu_vecs_gathered[2],
                  gpu_vecs_gathered[0],
                  gpu_vecs_gathered[1],
              ],
              ["fc"],
          )
          _, loss = model.net.SoftmaxWithLoss(
              [fc, "label"],
              ["ce_loss", "avg_loss"],
              only_loss=True,
          )
          loss = model.Scale(loss, scale=loss_scale)
          model.net.Print(loss, [], limit=10)
          return [loss]
  
      def param_update_fun(model):
          ONE = model.param_init_net.ConstantFill(
              [],
              "ONE",
              shape=[1],
              value=1.0,
          )
          LR = model.CopyCPUToGPU(self.LR, "LR")
          for param in model.GetParams():
              param_grad = model.param_to_grad[param]
              if not isinstance(param_grad, core.GradientSlice):
                  model.WeightedSum([param, ONE, param_grad, LR], param)
              else:
                  model.net.ScatterWeightedSum(
                      [
                          param,
                          ONE,
                          param_grad.indices,
                          param_grad.values,
                          ONE,
                      ],
                      param,
                  )
  
      workspace.ResetWorkspace()
      model = cnn.CNNModelHelper(
          order="NHWC",
          name="sparse_test{}".format(gpu_devices),
      )
      batch_size = 32
      batch_per_device = batch_size // len(gpu_devices)
  
      with core.NameScope("cpu"):
          with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU)):
              self.ITER = model.Iter("ITER")
              self.LR = model.net.LearningRate(
                  [self.ITER],
                  "LR",
                  base_lr=(-0.1),
                  policy="fixed",
              )
              """
                  self.vecs consists of 3 big blobs on which we call Gather:
                  1) FC weights, shape=(V, 16)
                  2) FC bias, shape=(V)
                  3) FC input, shape=(batch_per_device, 16)
                  """
              self.vecs = [
                  model.param_init_net.UniformFill(
                      [], "vec_{}".format(num), shape=[V, 16]
                  )
                  for num in range(2)
              ]
              self.vecs.append(
                  model.param_init_net.UniformFill(
                      [], "vec_2", shape=[batch_per_device, 16]
                  )
              )
              self.ONE_CPU = model.param_init_net.ConstantFill(
                  [],
                  "ONE_CPU",
                  shape=[1],
                  value=1.0,
              )
  
      data_parallel_model.Parallelize_GPU(
          model,
          input_builder_fun=input_builder_fun,
          forward_pass_builder_fun=model_build_fun,
          param_update_builder_fun=param_update_fun,
          devices=gpu_devices,
      )
  
      # Update the vecs
      with core.DeviceScope(
          core.DeviceOption(workspace.GpuDeviceType, 0)
      ):
          for num, vec in enumerate(self.vecs[:-1]):
              model.CopyGPUToCPU("gpu_0/gpuvec_{}".format(num), vec)
  
      # Each run has same input, independent of number of gpus
      for i in range(0, 10):
          np.random.seed(2603)
          full_indices = np.random.permutation(V)[:batch_size].reshape(
              batch_size
          )
          full_labels = full_indices[:] % batch_per_device
  
          for j, g in enumerate(gpu_devices):
              st = j * batch_per_device
              en = st + batch_per_device
              indices = full_indices[st:en].astype(np.int32)
              labels = full_labels[st:en].astype(np.int32)
  
              with core.DeviceScope(
                  core.DeviceOption(workspace.GpuDeviceType, g)
              ):
                  workspace.FeedBlob(
                      "gpu_{}/indices".format(g), indices
                  )
                  workspace.FeedBlob("gpu_{}/label".format(g), labels)
  
          if i == 0:
              workspace.RunNetOnce(model.param_init_net)
              # Force vecs to be same on all runs
              orig_vecs = [
                  np.random.rand(V, 16).astype(np.float32),
                  np.random.rand(V).astype(np.float32),
                  np.random.rand(V, 16).astype(np.float32),
              ]
              for vec, orig_vec in zip(self.vecs, orig_vecs):
                  workspace.FeedBlob(vec, orig_vec)
              for g in gpu_devices:
                  for num, orig_vec in enumerate(orig_vecs):
                      workspace.FeedBlob(
                          "gpu_{}/gpuvec_{}".format(g, num),
                          orig_vec,
                          device_option=core.DeviceOption(
                              workspace.GpuDeviceType, g
                          ),
                      )
              workspace.CreateNet(model.net)
  
          workspace.RunNet(model.net.Proto().name)
  
          idx = workspace.FetchBlob("gpu_0/indices")
          grad_slices = [
              workspace.FetchBlob(
                  "gpu_{}/gpu_vec_gathered_{}_grad".format(g, num)
              )
              for g in gpu_devices
              for num in range(2)
          ]
          for grad_slice in grad_slices:
              # print (len(idx), len(grad_slice))
              assert len(idx) == len(grad_slice), (
                  "Number of indices {} is not same as number of gradient "
                  "slices {}. This might lead to illegal memory access".format(
                      len(idx), len(grad_slice)
                  )
              )
  
imports: from multiprocessing import Process, Queue
  import numpy
  import os
  import shutil
  import tempfile
  import unittest
  import time
  from unittest.mock import Mock
  from hypothesis import assume, given, settings
  import hypothesis.strategies
  from caffe2.proto import caffe2_pb2
  from caffe2.python import brew, core, cnn, data_parallel_model, dyndep, model_helper, optimizer, rnn_cell, workspace
  from caffe2.python.test_util import TestCase
label_desc: Use image node_id 1 for calling the SparseDataParallelModelTestWithSharedIndices obj's underlying member method code with example usage: obj.run_model(V, gpu_devices) and returns: None, list