code: def test_compress_kvstore(kv_type, compression="2bit", threshold=0.5):
      print(
          kv_type
          + " with "
          + compression
          + " compression and threshold is "
          + str(threshold)
      )
      rate = 2
      quantize_func = None
      if compression == "1bit":
          quantize_func = compute_1bit
      elif compression == "2bit":
          quantize_func = compute_2bit
      else:
          raise RuntimeError("Unknown gradient compression type!")
      kv = mx.kv.create(kv_type)
      kv.set_gradient_compression(
          {"type": compression, "threshold": threshold}
      )
      kv.set_optimizer(mx.optimizer.create("test", learning_rate=-rate))
      for k, s in zip(keys, shapes):
          kv.init(k, mx.nd.zeros(s))
      # init one key with 1s so we can check if it was compressed during init
      kv.init(gc_init_test_key, mx.nd.ones(shapes[0]))
      # use different keys for random tests so that
      # we can track residual from start
      random_keys = [13, 15, 17]
      for k, s in zip(random_keys, shapes):
          kv.init(k, mx.nd.zeros(s))
  
      def pull_init_test(kv):
          # checks that compression is not applied to init of key
          out = [
              mx.nd.zeros(shapes[0], mx.gpu(g)) for g in range(nworker)
          ]
          kv.pull(gc_init_test_key, out=out)
          exp = np.ones_like(out[0].asnumpy())
          for o in out:
              assert_almost_equal(o.asnumpy(), exp)
  
      def pull_before_push(kv):
          for _ in range(nrepeat):
              for j in range(len(keys)):
                  out = [
                      mx.nd.ones(shapes[j], mx.gpu(g))
                      for g in range(nworker)
                  ]
                  kv.pull(keys[j], out=out)
                  exp = np.zeros_like(out[0].asnumpy())
                  for o in out:
                      assert_almost_equal(o.asnumpy(), exp)
  
      def push_ones(kv, sign=1):
          for i in range(nrepeat):
              for j in range(len(keys)):
                  kv.push(
                      keys[j],
                      [
                          sign * mx.nd.ones(shapes[j], mx.gpu(g))
                          for g in range(nworker)
                      ],
                  )
                  out = [
                      mx.nd.zeros(shapes[j], mx.gpu(g))
                      for g in range(nworker)
                  ]
                  kv.pull(keys[j], out=out)
                  if sign == 1:
                      exp = (
                          (i + 1)
                          * rate
                          * nworker
                          * np.ones_like(out[0].asnumpy())
                      )
                  else:
                      exp = (
                          (nrepeat - i - 1)
                          * rate
                          * nworker
                          * np.ones_like(out[0].asnumpy())
                      )
                  for o in out:
                      assert_almost_equal(o.asnumpy(), exp)
  
      def verify_residual_1bit(kv, threshold, rate):
          # current values must equal to zero
          for j in range(len(keys)):
              out = [
                  mx.nd.ones(shapes[j], mx.gpu(g))
                  for g in range(nworker)
              ]
              kv.pull(keys[j], out=out)
              exp = np.zeros_like(out[0].asnumpy())
              for o in out:
                  assert_almost_equal(o.asnumpy(), exp)
  
          curr_residual = 0
          curr_val = (
              rate * nworker if 2 > threshold else -rate * nworker
          )
          for j in range(len(keys)):
              kv.push(
                  keys[j],
                  [
                      2 * mx.nd.ones(shapes[j], mx.gpu(g))
                      for g in range(nworker)
                  ],
              )
              out = [
                  mx.nd.zeros(shapes[j], mx.gpu(g))
                  for g in range(nworker)
              ]
              kv.pull(keys[j], out=out)
  
              for o in out:
                  check_diff_to_scalar(o, curr_val)
  
          curr_residual = 1 if 2 > threshold else 3
          curr_val += (
              rate * nworker
              if 0 + curr_residual > threshold
              else -rate * nworker
          )
          for j in range(len(keys)):
              kv.push(
                  keys[j],
                  [
                      mx.nd.zeros(shapes[j], mx.gpu(g))
                      for g in range(nworker)
                  ],
              )
              out = [
                  mx.nd.zeros(shapes[j], mx.gpu(g))
                  for g in range(nworker)
              ]
              kv.pull(keys[j], out=out)
              for o in out:
                  check_diff_to_scalar(o, curr_val)
  
          curr_residual += -1 if curr_residual > threshold else +1
          curr_val += (
              rate * nworker
              if -2 + curr_residual > threshold
              else -rate * nworker
          )
          for j in range(len(keys)):
              kv.push(
                  keys[j],
                  [
                      -2 * mx.nd.ones(shapes[j], mx.gpu(g))
                      for g in range(nworker)
                  ],
              )
              out = [
                  mx.nd.zeros(shapes[j], mx.gpu(g))
                  for g in range(nworker)
              ]
              kv.pull(keys[j], out=out)
              for o in out:
                  check_diff_to_scalar(o, curr_val)
  
      def push_zeros(kv):
          for _ in range(nrepeat):
              for j in range(len(keys)):
                  kv.push(
                      keys[j],
                      [
                          mx.nd.zeros(shapes[j], mx.gpu(g))
                          for g in range(nworker)
                      ],
                  )
                  out = [
                      mx.nd.ones(shapes[j], mx.gpu(g))
                      for g in range(nworker)
                  ]
                  kv.pull(keys[j], out=out)
                  exp = np.zeros_like(out[0].asnumpy())
                  for o in out:
                      assert_almost_equal(o.asnumpy(), exp)
  
      def verify_residual_2bit(kv, threshold, rate):
          for j in range(len(keys)):
              kv.push(
                  keys[j],
                  [
                      mx.nd.ones(shapes[j], mx.gpu(g)) * 0.4
                      for g in range(nworker)
                  ],
              )
              out = [
                  mx.nd.zeros(shapes[j], mx.gpu(g))
                  for g in range(nworker)
              ]
              kv.pull(keys[j], out=out)
              for o in out:
                  check_diff_to_scalar(o, 0)
  
              kv.push(
                  keys[j],
                  [
                      mx.nd.ones(shapes[j], mx.gpu(g))
                      * (threshold - 0.3)
                      for g in range(nworker)
                  ],
              )
              out = [
                  mx.nd.zeros(shapes[j], mx.gpu(g))
                  for g in range(nworker)
              ]
              kv.pull(keys[j], out=out)
              curval = threshold * rate * nworker
              for o in out:
                  check_diff_to_scalar(o, curval)
  
              kv.push(
                  keys[j],
                  [
                      mx.nd.ones(shapes[j], mx.gpu(g)) * (0.2)
                      for g in range(nworker)
                  ],
              )
              out = [
                  mx.nd.zeros(shapes[j], mx.gpu(g))
                  for g in range(nworker)
              ]
              kv.pull(keys[j], out=out)
              for o in out:
                  check_diff_to_scalar(o, curval)
  
              kv.push(
                  keys[j],
                  [
                      mx.nd.ones(shapes[j], mx.gpu(g))
                      * (threshold - 0.3)
                      for g in range(nworker)
                  ],
              )
              out = [
                  mx.nd.zeros(shapes[j], mx.gpu(g))
                  for g in range(nworker)
              ]
              kv.pull(keys[j], out=out)
              curval += threshold * rate * nworker
              for o in out:
                  check_diff_to_scalar(o, curval)
              # residual would be 0 now
          return curval
  
      def check_neg(kv, neg, rate, curval):
          for _ in range(nrepeat):
              curval = curval + rate * nworker * neg
              for j in range(len(keys)):
                  kv.push(
                      keys[j],
                      [
                          mx.nd.ones(shapes[j], mx.gpu(g)) * neg
                          for g in range(nworker)
                      ],
                  )
                  out = [
                      mx.nd.ones(shapes[j], mx.gpu(g))
                      for g in range(nworker)
                  ]
                  kv.pull(keys[j], out=out)
                  for o in out:
                      check_diff_to_scalar(o, curval)
              # residual would be 0 again
  
      def check_compr_random(kv, threshold):
          for k, s in zip(random_keys, shapes):
              curr_residual = [np.zeros(s) for g in range(nworker)]
              orig_val = [
                  mx.nd.zeros(s, mx.gpu(g)) for g in range(nworker)
              ]
              kv.pull(k, out=orig_val)
              grads = [
                  mx.nd.random_uniform(
                      -0.6, 0.6, shape=s, ctx=mx.gpu(g)
                  )
                  for g in range(nworker)
              ]
              grads_cpy = copy.deepcopy(grads)
              kv.push(k, grads)
              val = [mx.nd.zeros(s, mx.gpu(g)) for g in range(nworker)]
              kv.pull(k, out=val)
              diffs = [val[g] - orig_val[g] for g in range(nworker)]
              # compute expected by using simulation of operator
              # on cpu
              sum_dequantized_vals = np.zeros(s)
              for g in range(nworker):
                  (
                      compr,
                      curr_residual[g],
                      decompr,
                  ) = compute_expected_quantization(
                      grads_cpy[g],
                      curr_residual[g],
                      threshold,
                      quantize_func,
                  )
                  sum_dequantized_vals += decompr * rate
  
              for g in range(nworker):
                  assert_almost_equal(
                      diffs[g].asnumpy(), sum_dequantized_vals
                  )
  
      pull_init_test(kv)
      pull_before_push(kv)
      if compression == "1bit":
          push_ones(kv, sign=1)
          push_ones(kv, sign=-1)
          verify_residual_1bit(kv, threshold, rate)
      elif compression == "2bit":
          push_zeros(kv)
          curval = verify_residual_2bit(kv, threshold, rate)
          check_neg(kv, -1 * threshold, rate, curval)
      check_compr_random(kv, threshold)
  
file_path: mxnet/tests/nightly/test_kvstore.py
imports: import sys
  import mxnet
  import numpy
  import numpy.random
  import copy
  from mxnet.test_utils import assert_almost_equal
label_desc: Use image node_id 6 for calling a global function with example usage: test_compress_kvstore(kv_type, compression, threshold) and returns: curval