import tensorflow as tf

  def test_prob_and_grad_gives_finite_results_for_common_events(self):
    with self.test_session():
      mu = tf.Variable(0.0, name="mu")
      sigma = tf.Variable(1.0, name="sigma")
      qdist = distributions.QuantizedDistribution(
          base_dist_cls=distributions.Normal,
          mu=mu,
          sigma=sigma)
      x = tf.ceil(4 * self._rng.rand(100).astype(np.float32) - 2)

      tf.initialize_all_variables().run()

      proba = qdist.prob(x)
      self._assert_all_finite(proba.eval())

      grads = tf.gradients(proba, [mu, sigma])
      self._assert_all_finite(grads[0].eval())
      self._assert_all_finite(grads[1].eval())

  def test_lower_cutoff_must_be_below_upper_cutoff_or_we_raise(self):
    with self.test_session():
      qdist = distributions.QuantizedDistribution(
          base_dist_cls=distributions.Normal,
          lower_cutoff=1.,  # not strictly less than upper_cutoff.
          upper_cutoff=1.,
          mu=0.,
          sigma=1.,
          validate_args=True)

      self.assertTrue(qdist.validate_args)  # Default is True.
