code: def test_device_pushpull():
      def check_dense_pushpull(kv_type):
          for shape, key in zip(shapes, keys):
              for n_gpus in gpus:
                  kv_device = mx.kv.create(kv_type)
                  a = mx.nd.ones(shape, mx.gpu(0))
                  cur_key = str(key * max(gpus) + n_gpus)
                  kv_device.init(cur_key, a)
                  arr_list = [
                      mx.nd.ones(shape, mx.gpu(x))
                      for x in range(n_gpus)
                  ]
                  res = [
                      mx.nd.zeros(shape, mx.gpu(x))
                      for x in range(n_gpus)
                  ]
                  kv_device.push(cur_key, arr_list)
                  kv_device.pull(cur_key, res)
                  for x in range(n_gpus):
                      assert (
                          np.sum(np.abs((res[x] - n_gpus).asnumpy()))
                          == 0
                      )
  
      kvstore_tree_array_bound_values = [None, "1"]
      kvstore_usetree_values = [None, "1"]
      for y in kvstore_tree_array_bound_values:
          for x in kvstore_usetree_values:
              with environment(
                  {
                      "MXNET_KVSTORE_USETREE": x,
                      "MXNET_KVSTORE_TREE_ARRAY_BOUND": y,
                  }
              ):
                  check_dense_pushpull("local")
                  check_dense_pushpull("device")
  
imports: import mxnet
  import numpy
  import pytest
  import os
  import logging
  from mxnet.test_utils import environment
label_desc: Use image node_id 1 for calling a global function with example usage: test_device_pushpull() without return types