code: def test_quantization_bits_stochastic_rounding(self, bits, min_max):
      stage = stages_impl.UniformQuantizationEncodingStage(
          bits=bits, min_max=min_max, stochastic=True
      )
      test_data = self.run_one_to_many_encode_decode(
          stage, self.default_input
      )
      self._assert_is_integer_float(
          test_data.encoded_x[
              stages_impl.UniformQuantizationEncodingStage.ENCODED_VALUES_KEY
          ]
      )
      # For stochastic rounding, the potential error incurred by quantization
      # is bounded by the range of the input values divided by the number of
      # quantization buckets.
      if min_max is None:
          self.assertAllClose(
              test_data.x,
              test_data.decoded_x,
              rtol=0.0,
              atol=2 / (2**bits - 1),
          )
      else:
          self.assertAllClose(
              np.clip(test_data.x, -0.5, 0.5),
              test_data.decoded_x,
              rtol=0.0,
              atol=1 / (2**bits - 1),
          )
  
imports: import itertools
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tensorflow_model_optimization.python.core.internal.tensor_encoding.stages import stages_impl
  from tensorflow_model_optimization.python.core.internal.tensor_encoding.testing import test_utils
label_desc: Use image node_id 6 for calling the UniformQuantizationEncodingStageStageTest obj's underlying member method code with example usage: obj.test_quantization_bits_stochastic_rounding(bits, min_max) without return types