code: def test_sync_push_pull(nrepeat):
      def check_default_keys(dtype, nrepeat):
          # checks pull after push in loop, because behavior during
          # consecutive pushes doesn't offer any guarantees
          ks = keys_shapes if dtype == "float32" else fp16_keys_shapes
          for k, s in ks:
              for i in range(nrepeat):
                  kv.push(k, mx.nd.ones(s, dtype=dtype) * (my_rank + 1))
                  num = (nworker + 1) * nworker * rate / 2 * (i + 1) + 1
                  val = mx.nd.zeros(s, dtype=dtype)
                  kv.pull(k, out=val)
                  check_diff(val, num)
  
      def check_row_sparse_keys(dtype, nrepeat):
          # prepare gradient
          v = mx.nd.zeros(shape, dtype=dtype)
          my_row = my_rank % shape[0]
          v[my_row] = my_rank + 1
          # push
          if dtype == "float32":
              k = rsp_keys_shape[0]
          else:
              k = fp16_rsp_keys_shape[0]
          s = shape
          for i in range(nrepeat):
              kv.push(k, v.tostype("row_sparse"))
              # select a random subset of rows this worker is interested in
              num_rows = s[0]
              row_ids_np = np.random.randint(num_rows, size=num_rows)
              row_ids = (
                  mx.nd.array(row_ids_np)
                  .reshape((num_rows / 2, 2))
                  .astype(dtype)
              )
              # perform pull
              val = mx.nd.zeros(s, stype="row_sparse", dtype=dtype)
              kv.row_sparse_pull(k, out=val, row_ids=row_ids)
              # prepare updated values
              updated_val = mx.nd.ones(s, dtype=dtype)
              for rank in range(nworker):
                  row = rank % s[0]
                  updated_val[row] += (rank + 1) * rate * (i + 1)
              # verify subset of updated values
              expected = mx.nd.zeros(s, dtype=dtype)
              for row in row_ids_np:
                  expected[row] = updated_val[row]
              check_diff(val, expected, kv.rank)
  
      def check_row_sparse_keys_with_zeros(dtype, nrepeat):
          if dtype == "float32":
              k1 = rsp_keys_shape[1]
              k2 = rsp_keys_big_shape[0]
          else:
              k1 = fp16_rsp_keys_shape[1]
              k2 = fp16_rsp_keys_big_shape[0]
          # prepare gradient
          v = mx.nd.sparse.zeros("row_sparse", shape, dtype=dtype)
          big_v = mx.nd.sparse.zeros(
              "row_sparse", big_shape, dtype=dtype
          )
          # push
          for _ in range(nrepeat):
              kv.push(k1, v)
              kv.push(k2, big_v)
              # pull a subset of rows this worker is interested in
              all_row_ids = np.arange(shape[0])
              val = mx.nd.sparse.zeros("row_sparse", shape)
              big_val = mx.nd.sparse.zeros("row_sparse", big_shape)
              kv.row_sparse_pull(
                  k1, out=val, row_ids=mx.nd.array(all_row_ids)
              )
              big_all_row_ids = np.arange(big_shape[0])
              kv.row_sparse_pull(
                  k2, out=big_val, row_ids=mx.nd.array(big_all_row_ids)
              )
              # verify results
              check_diff(val, 1)
              check_diff(big_val, 1)
              # pull empty weights
              kv.row_sparse_pull(k1, out=val, row_ids=mx.nd.array([]))
              kv.row_sparse_pull(
                  k2, out=big_val, row_ids=mx.nd.array([])
              )
              check_diff(val, 0)
              check_diff(big_val, 0)
  
      def check_big_row_sparse_keys(dtype, nrepeat):
          if dtype == "float32":
              k = rsp_keys_big_shape[0]
          else:
              k = fp16_rsp_keys_big_shape[0]
  
          mx.random.seed(123)
          rnd.seed(123)
          density = 0.3
          # prepare gradient
          v = mx.nd.zeros(big_shape, dtype=dtype)
          idx_sample = rnd.rand(big_shape[0])
          indices = np.argwhere(idx_sample < density).flatten()
          # each worker chooses a subset of the indices to update
          update_rows = []
          for rank in range(nworker):
              rows = []
              i = 0
              step = (rank + 1) * 2
              while i < len(indices):
                  rows.append(indices[i])
                  i += step
              update_rows.append(np.array(rows))
          # rows to update for this worker
          for row in update_rows[my_rank]:
              v[row] = my_rank + 1
          # push
          for i in range(nrepeat):
              kv.push(k, v.tostype("row_sparse"))
              # select a random subset of rows this worker is interested in
              mx.random.seed(my_rank)
              rnd.seed(my_rank)
              num_rows = big_shape[0]
              row_ids_np = np.random.randint(num_rows, size=num_rows)
              row_ids = (
                  mx.nd.array(row_ids_np)
                  .reshape((num_rows / 2, 2))
                  .astype(dtype)
              )
              # perform pull
              val = mx.nd.zeros(
                  big_shape, stype="row_sparse", dtype=dtype
              )
              kv.row_sparse_pull(k, out=val, row_ids=row_ids)
              # prepare expected result
              updated_val = mx.nd.ones(big_shape, dtype=dtype)
              # apply updates from each worker
              for rank in range(nworker):
                  for row in update_rows[rank]:
                      updated_val[row] += (rank + 1) * rate * (i + 1)
  
              expected = mx.nd.zeros(big_shape, dtype=dtype)
              for row in row_ids_np:
                  expected[row] = updated_val[row]
              check_diff(val, expected.astype(dtype), rank=my_rank)
  
      for dtype in ["float16", "float32"]:
          check_default_keys(dtype, nrepeat)
          check_row_sparse_keys(dtype, nrepeat)
          check_row_sparse_keys_with_zeros(dtype, nrepeat)
          check_big_row_sparse_keys(dtype, nrepeat)
      print(
          "worker "
          + str(my_rank)
          + " is done with non compression tests"
      )
  
file_path: mxnet/tests/nightly/dist_sync_kvstore.py
imports: import sys
  import argparse
  import mxnet
  import numpy
  import numpy.random
  from mxnet.test_utils import assert_almost_equal, assert_exception
  from test_kvstore import compute_expected_quantization, compute_1bit, compute_2bit
label_desc: Use image node_id 5 for calling a global function with example usage: test_sync_push_pull(nrepeat) without return types