import tensorflow as tf

  def test_quantization_of_batch_of_uniforms(self):
    batch_shape = (5, 5)
    with self.test_session():
      # The uniforms are supported on [0, 10].  The qdist considers the
      # intervals
      # ... (0, 1](1, 2]...(9, 10]...
      # with the intervals displayed above each holding 1 / 10 of the mass.
      # The qdist will be defined with no cutoffs,
      qdist = distributions.QuantizedDistribution(
          base_dist_cls=distributions.Uniform,
          lower_cutoff=None,
          upper_cutoff=None,
          a=tf.zeros(
              batch_shape, dtype=tf.float32),
          b=10 * tf.ones(
              batch_shape, dtype=tf.float32))

      # x is random integers in {-3,...,12}.
      x = self._rng.randint(-3, 13, size=batch_shape).astype(np.float32)

      # pmf
      # qdist.pmf(j) = 1 / 10 for j in {1,...,10}, and 0 otherwise,
      expected_pmf = (1 / 10) * np.ones(batch_shape)
      expected_pmf[x < 1] = 0.
