code: def uniform_binning_correction(x, n_bits=8):
    
    n_bins = 2**n_bits
    batch_size, height, width, n_channels = common_layers.shape_list(x)
    hwc = float(height * width * n_channels)
  
    x = x + tf.random_uniform(
        shape=(batch_size, height, width, n_channels),
        minval=0.0, maxval=1.0/n_bins)
    objective = -np.log(n_bins) * hwc * tf.ones(batch_size)
    return x, objective
docs: Replaces x^i with q^i(x) = U(x, x + 1.0 / 256.0).
  
    Args:
      x: 4-D Tensor of shape (NHWC)
      n_bits: optional.
    Returns:
      x: x ~ U(x, x + 1.0 / 256)
      objective: Equivalent to -q(x)*log(q(x)).