code: def dithering(x, k, state, partition="linear", norm="max"):
      y = x.flatten()
      if norm == "max":
          scale = np.max(np.abs(y))
      elif norm == "l2":
          scale = np.linalg.norm(y.astype(np.float64), ord=2)
      else:
          raise ValueError("Unsupported normalization")
      y /= scale
      sign = np.sign(y)
      y = np.abs(y)
  
      # stocastic rounding
      if partition == "linear":
          y *= k
          low = np.floor(y)
          p = y - low  # whether to ceil
          y = low + bernoulli(p, state)
          y /= k
      elif partition == "natural":
          y *= 2 ** (k - 1)
          low = round_next_pow2((np.ceil(y).astype(np.uint32))) >> 1
          length = copy.deepcopy(low)
          length[length == 0] = 1
          p = (y - low) / length
          y = low + length * bernoulli(p, state)
          y = y.astype(np.float32)
          y /= 2 ** (k - 1)
      else:
          raise ValueError("Unsupported partition")
  
      y *= sign
      y *= scale
      return y.reshape(x.shape)
  
file_path: byteps/tests/test_dithering.py
imports: import copy
  import itertools
  import unittest
  import byteps.mxnet
  import mxnet
  import mxnet.ndarray
  import numpy
  from gluoncv.model_zoo import get_model
  from mxnet import autograd, gluon
  from numba import jit
  from parameterized import parameterized
  from tqdm import tqdm
  from meta_test import MetaTest
  from utils import bernoulli, fake_data
label_desc: Use image node_id 2 for calling a global function with example usage: dithering(x, k, state, partition, norm) and returns: y