code: def test_sync_2bit_compression(threshold, nrepeat):
      def check_compr_residual(threshold):
          for k, s in compr_keys_shapes:
              # doesn't meet threshold
              kv.push(k, mx.nd.ones(s) * 0.4)
              val = mx.nd.zeros(s)
              kv.pull(k, val)
              check_diff(val, 0)
  
              # just meets threshold with residual
              kv.push(k, mx.nd.ones(s) * (threshold - 0.4))
              val2 = mx.nd.zeros(s)
              kv.pull(k, val2)
              curval = threshold * rate * nworker
              check_diff(val2, curval)
  
              # doesn't meet threshold
              kv.push(k, mx.nd.ones(s) * 0.2)
              val3 = mx.nd.zeros(s)
              kv.pull(k, val3)
              check_diff(val3, curval)
  
              # exceeds again
              kv.push(k, mx.nd.ones(s) * (threshold - 0.2))
              val4 = mx.nd.zeros(s)
              kv.pull(k, val4)
              curval += threshold * rate * nworker
              check_diff(val4, curval)
              # residual is 0 now
  
      def check_compr_ones(threshold):
          for k, s in compr_keys_shapes:
              val = mx.nd.zeros(s)
              kv.pull(k, val)
              curval = val[0][0].asnumpy()[0]
              kv.push(k, mx.nd.ones(s) * threshold)
              val2 = mx.nd.zeros(s)
              kv.pull(k, val2)
              newval = curval + rate * nworker * threshold
              check_diff(val2, newval)
              # residual = 0  again
  
      def check_compr_pull_before_push():
          for k, s in compr_keys_shapes:
              val = mx.nd.ones(s)
              kv.pull(k, val)
              check_diff(val, 0)
          for k, s in compr_init_keys_shapes:
              # tests that GC is not used for init of a key
              val = mx.nd.zeros(s)
              kv.pull(k, val)
              check_diff(val, 1)
  
      def check_compr_zero():
          for k, s in compr_keys_shapes:
              kv.push(k, mx.nd.zeros(s))
              # to check that all are set to 0s
              val = mx.nd.ones(s)
              kv.pull(k, val)
              check_diff(val, 0)
  
      def check_compr_random(threshold, nrepeat):
          # set a seed so all workers generate same data. knowing this helps
          # calculate expected value after pull
          mx.random.seed(123)
          rnd.seed(123)
  
          # use new keys so residual is 0 for calculation of expected
          for k, s in compr_random_keys_shapes:
              kv.init(k, mx.nd.zeros(s))
          for k, s in compr_random_keys_shapes:
              curr_residual = np.zeros(s)
              for _ in range(nrepeat):
                  orig_val = mx.nd.zeros(s)
                  kv.pull(k, orig_val)
  
                  grad = mx.nd.array(rnd.rand(s[0], s[1]))
                  # creates a copy because push changes grad because of assignment
                  grad_cpy = mx.nd.array(grad)
                  kv.push(k, grad)
                  val = mx.nd.zeros(s)
                  kv.pull(k, val)
  
                  diff = val - orig_val
  
                  # compute expected by using simulation of operator
                  (
                      compr,
                      curr_residual,
                      decompr,
                  ) = compute_expected_quantization(
                      grad_cpy, curr_residual, threshold, compute_2bit
                  )
                  decompr *= nworker * rate
                  assert_almost_equal(diff.asnumpy(), decompr)
  
      print(
          "worker "
          + str(my_rank)
          + " started with 2bit compression tests"
      )
      check_compr_pull_before_push()
      check_compr_zero()
      check_compr_residual(threshold)
      check_compr_ones(threshold)
      check_compr_random(threshold, nrepeat)
      print(
          "worker "
          + str(my_rank)
          + " is done with 2bit 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 7 for calling a global function with example usage: test_sync_2bit_compression(threshold, nrepeat) without return types