text: # Copyright 2018 The TensorFlow Probability Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ============================================================================
  """Tests for TransformedDistribution."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  # Dependency imports
  import numpy as np
  from scipy import stats
  
  import tensorflow.compat.v1 as tf1
  import tensorflow.compat.v2 as tf
  import tensorflow_probability as tfp
  
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util as tfp_test_util
  from tensorflow.python.framework import test_util  # pylint: disable=g-direct-tensorflow-import
  
  tfb = tfp.bijectors
  tfd = tfp.distributions
  
  
  class DummyMatrixTransform(tfb.Bijector):
    """Tractable matrix transformation.
  
    This is a non-sensical bijector that has forward/inverse_min_event_ndims=2.
    The main use is to check that transformed distribution calculations are done
    appropriately.
    """
  
    def __init__(self):
      super(DummyMatrixTransform, self).__init__(
          forward_min_event_ndims=2,
          is_constant_jacobian=False,
          validate_args=False,
          name="dummy")
  
    def _forward(self, x):
      return x
  
    def _inverse(self, y):
      return y
  
    # Note: These jacobians don't make sense.
    def _forward_log_det_jacobian(self, x):
      return -tf.linalg.det(x)
  
    def _inverse_log_det_jacobian(self, x):
      return tf.linalg.det(x)
  
  
  @test_util.run_all_in_graph_and_eager_modes
  class TransformedDistributionTest(tf.test.TestCase):
  
    def _cls(self):
      return tfd.TransformedDistribution
  
    def _make_unimplemented(self, name):
      def _unimplemented(self, *args):  # pylint: disable=unused-argument
        raise NotImplementedError("{} not implemented".format(name))
      return _unimplemented
  
    def testTransformedDistribution(self):
      mu = 3.0
      sigma = 2.0
      # Note: the Jacobian callable only works for this example; more generally
      # you may or may not need a reduce_sum.
      log_normal = self._cls()(
          distribution=tfd.Normal(loc=mu, scale=sigma), bijector=tfb.Exp())
      sp_dist = stats.lognorm(s=sigma, scale=np.exp(mu))
  
      # sample
      sample = log_normal.sample(100000, seed=tfp_test_util.test_seed())
      self.assertAllEqual([], log_normal.event_shape)
      self.assertAllEqual([], self.evaluate(log_normal.event_shape_tensor()))
      self.assertAllClose(
          sp_dist.mean(), np.mean(self.evaluate(sample)), atol=0.0, rtol=0.05)
  
      # pdf, log_pdf, cdf, etc...
      # The mean of the lognormal is around 148.
      test_vals = np.linspace(0.1, 1000., num=20).astype(np.float32)
      for func in [[log_normal.log_prob, sp_dist.logpdf],
                   [log_normal.prob, sp_dist.pdf],
                   [log_normal.log_cdf, sp_dist.logcdf],
                   [log_normal.cdf, sp_dist.cdf],
                   [log_normal.survival_function, sp_dist.sf],
                   [log_normal.log_survival_function, sp_dist.logsf]]:
        actual = func[0](test_vals)
        expected = func[1](test_vals)
        self.assertAllClose(
            expected, self.evaluate(actual), atol=0, rtol=0.01)
  
    def testNonInjectiveTransformedDistribution(self):
      mu = 1.
      sigma = 2.0
      abs_normal = self._cls()(
          distribution=tfd.Normal(loc=mu, scale=sigma),
          bijector=tfb.AbsoluteValue())
      sp_normal = stats.norm(mu, sigma)
  
      # sample
      sample = abs_normal.sample(100000, seed=tfp_test_util.test_seed())
      self.assertAllEqual([], abs_normal.event_shape)
      sample_ = self.evaluate(sample)
      self.assertAllEqual([], self.evaluate(abs_normal.event_shape_tensor()))
  
      # Abs > 0, duh!
      np.testing.assert_array_less(0, sample_)
  
      # Let X ~ Normal(mu, sigma), Y := |X|, then
      # P[Y < 0.77] = P[-0.77 < X < 0.77]
      self.assertAllClose(
          sp_normal.cdf(0.77) - sp_normal.cdf(-0.77),
          (sample_ < 0.77).mean(), rtol=0.01)
  
      # p_Y(y) = p_X(-y) + p_X(y),
      self.assertAllClose(
          sp_normal.pdf(1.13) + sp_normal.pdf(-1.13),
          self.evaluate(abs_normal.prob(1.13)))
  
      # Log[p_Y(y)] = Log[p_X(-y) + p_X(y)]
      self.assertAllClose(
          np.log(sp_normal.pdf(2.13) + sp_normal.pdf(-2.13)),
          self.evaluate(abs_normal.log_prob(2.13)))
  
    def testQuantile(self):
      logit_normal = self._cls()(
          distribution=tfd.Normal(loc=0., scale=1.),
          bijector=tfb.Sigmoid(),
          validate_args=True)
      grid = [0., 0.25, 0.5, 0.75, 1.]
      q = logit_normal.quantile(grid)
      cdf = logit_normal.cdf(q)
      cdf_ = self.evaluate(cdf)
      self.assertAllClose(grid, cdf_, rtol=1e-6, atol=0.)
  
    def testCachedSamples(self):
      class ExpForwardOnly(tfb.Bijector):
  
        def __init__(self):
          super(ExpForwardOnly, self).__init__(forward_min_event_ndims=0)
  
        def _forward(self, x):
          return tf.exp(x)
  
        def _forward_log_det_jacobian(self, x):
          return tf.convert_to_tensor(value=x)
  
      exp_forward_only = ExpForwardOnly()
  
      mu = 3.0
      sigma = 0.02
      log_normal = self._cls()(
          distribution=tfd.Normal(loc=mu, scale=sigma), bijector=exp_forward_only)
  
      sample = log_normal.sample([2, 3], seed=tfp_test_util.test_seed())
      sample_val, log_pdf_val = self.evaluate(
          [sample, log_normal.log_prob(sample)])
      expected_log_pdf = stats.lognorm.logpdf(
          sample_val, s=sigma, scale=np.exp(mu))
      self.assertAllClose(expected_log_pdf, log_pdf_val, rtol=1e-4, atol=0.)
  
    def testCachedSamplesInvert(self):
      class ExpInverseOnly(tfb.Bijector):
  
        def __init__(self):
          super(ExpInverseOnly, self).__init__(inverse_min_event_ndims=0)
  
        def _inverse(self, y):
          return tf.math.log(y)
  
        def _inverse_log_det_jacobian(self, y):
          return -tf.math.log(y)
  
      exp_inverse_only = ExpInverseOnly()
  
      log_forward_only = tfb.Invert(exp_inverse_only)
  
      # The log bijector isn't defined over the whole real line, so we make
      # sigma sufficiently small so that the draws are positive.
      mu = 2.
      sigma = 1e-2
      exp_normal = self._cls()(
          distribution=tfd.Normal(loc=mu, scale=sigma), bijector=log_forward_only)
  
      sample = exp_normal.sample(
          [2, 3], seed=tfp_test_util.test_seed(hardcoded_seed=42))
      sample_val, log_pdf_val = self.evaluate(
          [sample, exp_normal.log_prob(sample)])
      expected_log_pdf = sample_val + stats.norm.logpdf(
          np.exp(sample_val), loc=mu, scale=sigma)
      self.assertAllClose(expected_log_pdf, log_pdf_val, atol=0., rtol=1e-5)
  
    def testShapeChangingBijector(self):
      softmax = tfb.SoftmaxCentered()
      standard_normal = tfd.Normal(loc=0., scale=1.)
      multi_logit_normal = self._cls()(
          distribution=standard_normal, bijector=softmax, event_shape=[1])
      x = [[[-np.log(3.)], [0.]], [[np.log(3)], [np.log(5)]]]
      x = np.float32(x)
      y = self.evaluate(softmax.forward(x))
      expected_log_pdf = (
          np.squeeze(stats.norm(loc=0., scale=1.).logpdf(x)) - np.sum(
              np.log(y), axis=-1))
      self.assertAllClose(expected_log_pdf,
                          self.evaluate(multi_logit_normal.log_prob(y)))
      self.assertAllClose(
          [1, 2, 3, 2],
          self.evaluate(tf.shape(input=multi_logit_normal.sample([1, 2, 3]))))
      self.assertAllEqual([2], multi_logit_normal.event_shape)
      self.assertAllEqual([2],
                          self.evaluate(multi_logit_normal.event_shape_tensor()))
  
    def testCastLogDetJacobian(self):
      """Test log_prob when Jacobian and log_prob dtypes do not match."""
  
      # Create an identity bijector whose jacobians have dtype int32
      int_identity = tfb.Inline(
          forward_fn=tf.identity,
          inverse_fn=tf.identity,
          inverse_log_det_jacobian_fn=(lambda y: tf.cast(0, tf.int32)),
          forward_log_det_jacobian_fn=(lambda x: tf.cast(0, tf.int32)),
          forward_min_event_ndims=0,
          is_constant_jacobian=True)
      normal = self._cls()(
          distribution=tfd.Normal(loc=0., scale=1.),
          bijector=int_identity,
          validate_args=True)
  
      y = normal.sample()
      self.evaluate(normal.log_prob(y))
      self.evaluate(normal.prob(y))
      self.evaluate(normal.mean())
      self.evaluate(normal.entropy())
  
    def testMean(self):
      shift = np.array([[-1, 0, 1], [-1, -2, -3]], dtype=np.float32)
      diag = np.array([[1, 2, 3], [2, 3, 2]], dtype=np.float32)
      fake_mvn = self._cls()(
          tfd.MultivariateNormalDiag(
              loc=tf.zeros_like(shift),
              scale_diag=tf.ones_like(diag),
              validate_args=True),
          tfb.AffineLinearOperator(
              shift,
              scale=tf.linalg.LinearOperatorDiag(diag, is_non_singular=True),
              validate_args=True),
          validate_args=True)
      self.assertAllClose(shift, self.evaluate(fake_mvn.mean()))
  
    def testMeanShapeOverride(self):
      shift = np.array([[-1, 0, 1], [-1, -2, -3]], dtype=np.float32)
      diag = np.array([[1, 2, 3], [2, 3, 2]], dtype=np.float32)
      fake_mvn = self._cls()(
          tfd.Normal(loc=0.0, scale=1.0),
          tfb.AffineLinearOperator(
              shift,
              scale=tf.linalg.LinearOperatorDiag(diag, is_non_singular=True),
              validate_args=True),
          batch_shape=[2],
          event_shape=[3],
          validate_args=True)
      self.assertAllClose(shift, self.evaluate(fake_mvn.mean()))
  
    def testEntropy(self):
      shift = np.array([[-1, 0, 1], [-1, -2, -3]], dtype=np.float32)
      diag = np.array([[1, 2, 3], [2, 3, 2]], dtype=np.float32)
      actual_mvn_entropy = np.concatenate(
          [[stats.multivariate_normal(shift[i], np.diag(diag[i]**2)).entropy()]
           for i in range(len(diag))])
      fake_mvn = self._cls()(
          tfd.MultivariateNormalDiag(
              loc=tf.zeros_like(shift),
              scale_diag=tf.ones_like(diag),
              validate_args=True),
          tfb.AffineLinearOperator(
              shift,
              scale=tf.linalg.LinearOperatorDiag(diag, is_non_singular=True),
              validate_args=True),
          validate_args=True)
      self.assertAllClose(actual_mvn_entropy, self.evaluate(fake_mvn.entropy()))
  
    def testScalarBatchScalarEventIdentityScale(self):
      exp2 = self._cls()(
          tfd.Exponential(rate=0.25), bijector=tfb.AffineScalar(scale=2.))
      log_prob = exp2.log_prob(1.)
      log_prob_ = self.evaluate(log_prob)
      base_log_prob = -0.5 * 0.25 + np.log(0.25)
      ildj = np.log(2.)
      self.assertAllClose(base_log_prob - ildj, log_prob_, rtol=1e-6, atol=0.)
  
  
  class ScalarToMultiTest(tf.test.TestCase):
  
    def _cls(self):
      return tfd.TransformedDistribution
  
    def setUp(self):
      self._shift = np.array([-1, 0, 1], dtype=np.float32)
      self._tril = np.array([[[1., 0, 0],
                              [2, 1, 0],
                              [3, 2, 1]],
                             [[2, 0, 0],
                              [3, 2, 0],
                              [4, 3, 2]]],
                            dtype=np.float32)
      super(ScalarToMultiTest, self).setUp()
  
    def _testMVN(self,
                 base_distribution_class,
                 base_distribution_kwargs,
                 batch_shape=(),
                 event_shape=(),
                 not_implemented_message=None):
      # Overriding shapes must be compatible w/bijector; most bijectors are
      # batch_shape agnostic and only care about event_ndims.
      # In the case of `Affine`, if we got it wrong then it would fire an
      # exception due to incompatible dimensions.
      batch_shape_pl = tf1.placeholder_with_default(
          input=np.int32(batch_shape), shape=None, name="dynamic_batch_shape")
      event_shape_pl = tf1.placeholder_with_default(
          input=np.int32(event_shape), shape=None, name="dynamic_event_shape")
      fake_mvn_dynamic = self._cls()(
          distribution=base_distribution_class(
              validate_args=True, **base_distribution_kwargs),
          bijector=tfb.Affine(shift=self._shift, scale_tril=self._tril),
          batch_shape=batch_shape_pl,
          event_shape=event_shape_pl,
          validate_args=True)
  
      fake_mvn_static = self._cls()(
          distribution=base_distribution_class(
              validate_args=True, **base_distribution_kwargs),
          bijector=tfb.Affine(shift=self._shift, scale_tril=self._tril),
          batch_shape=batch_shape,
          event_shape=event_shape,
          validate_args=True)
  
      actual_mean = np.tile(self._shift, [2, 1])  # Affine elided this tile.
      actual_cov = np.matmul(self._tril, np.transpose(self._tril, [0, 2, 1]))
  
      def actual_mvn_log_prob(x):
        return np.concatenate([[  # pylint: disable=g-complex-comprehension
            stats.multivariate_normal(actual_mean[i],
                                      actual_cov[i]).logpdf(x[:, i, :])
        ] for i in range(len(actual_cov))]).T
  
      actual_mvn_entropy = np.concatenate(
          [[stats.multivariate_normal(actual_mean[i], actual_cov[i]).entropy()]
           for i in range(len(actual_cov))])
  
      self.assertAllEqual([3], fake_mvn_static.event_shape)
      self.assertAllEqual([2], fake_mvn_static.batch_shape)
  
      if not tf.executing_eagerly():
        self.assertAllEqual(tf.TensorShape(None), fake_mvn_dynamic.event_shape)
        self.assertAllEqual(tf.TensorShape(None), fake_mvn_dynamic.batch_shape)
  
      x = self.evaluate(fake_mvn_static.sample(5, seed=tfp_test_util.test_seed()))
      for unsupported_fn in (fake_mvn_static.log_cdf, fake_mvn_static.cdf,
                             fake_mvn_static.survival_function,
                             fake_mvn_static.log_survival_function):
        with self.assertRaisesRegexp(NotImplementedError,
                                     not_implemented_message):
          unsupported_fn(x)
  
      num_samples = 7e3
      for fake_mvn in [fake_mvn_static, fake_mvn_dynamic]:
        # Ensure sample works by checking first, second moments.
        y = fake_mvn.sample(int(num_samples), seed=tfp_test_util.test_seed())
        x = y[0:5, ...]
        sample_mean = tf.reduce_mean(input_tensor=y, axis=0)
        centered_y = tf.transpose(a=y - sample_mean, perm=[1, 2, 0])
        sample_cov = tf.matmul(
            centered_y, centered_y, transpose_b=True) / num_samples
        [
            sample_mean_,
            sample_cov_,
            x_,
            fake_event_shape_,
            fake_batch_shape_,
            fake_log_prob_,
            fake_prob_,
            fake_mean_,
            fake_entropy_,
        ] = self.evaluate([
            sample_mean,
            sample_cov,
            x,
            fake_mvn.event_shape_tensor(),
            fake_mvn.batch_shape_tensor(),
            fake_mvn.log_prob(x),
            fake_mvn.prob(x),
            fake_mvn.mean(),
            fake_mvn.entropy(),
        ])
  
        self.assertAllClose(actual_mean, sample_mean_, atol=0.1, rtol=0.1)
        self.assertAllClose(actual_cov, sample_cov_, atol=0., rtol=0.1)
  
        # Ensure all other functions work as intended.
        self.assertAllEqual([5, 2, 3], x_.shape)
        self.assertAllEqual([3], fake_event_shape_)
        self.assertAllEqual([2], fake_batch_shape_)
        self.assertAllClose(
            actual_mvn_log_prob(x_), fake_log_prob_, atol=0., rtol=1e-6)
        self.assertAllClose(
            np.exp(actual_mvn_log_prob(x_)), fake_prob_, atol=0., rtol=1e-5)
        self.assertAllClose(actual_mean, fake_mean_, atol=0., rtol=1e-6)
        self.assertAllClose(actual_mvn_entropy, fake_entropy_, atol=0., rtol=1e-6)
  
    def testScalarBatchScalarEvent(self):
      self._testMVN(
          base_distribution_class=tfd.Normal,
          base_distribution_kwargs={
              "loc": 0.,
              "scale": 1.
          },
          batch_shape=[2],
          event_shape=[3],
          not_implemented_message="not implemented when overriding event_shape")
  
    def testScalarBatchNonScalarEvent(self):
      self._testMVN(
          base_distribution_class=tfd.MultivariateNormalDiag,
          base_distribution_kwargs={
              "loc": [0., 0., 0.],
              "scale_diag": [1., 1, 1]
          },
          batch_shape=[2],
          not_implemented_message="not implemented")
  
      # Can't override event_shape for scalar batch, non-scalar event.
      with self.assertRaisesRegexp(ValueError, "base distribution not scalar"):
        self._cls()(
            distribution=tfd.MultivariateNormalDiag(loc=[0.], scale_diag=[1.]),
            bijector=tfb.Affine(shift=self._shift, scale_tril=self._tril),
            batch_shape=[2],
            event_shape=[3],
            validate_args=True)
  
    def testNonScalarBatchScalarEvent(self):
      self._testMVN(
          base_distribution_class=tfd.Normal,
          base_distribution_kwargs={
              "loc": [0., 0],
              "scale": [1., 1]
          },
          event_shape=[3],
          not_implemented_message="not implemented when overriding event_shape")
  
      # Can't override batch_shape for non-scalar batch, scalar event.
      with self.assertRaisesRegexp(ValueError, "base distribution not scalar"):
        self._cls()(
            distribution=tfd.Normal(loc=[0.], scale=[1.]),
            bijector=tfb.Affine(shift=self._shift, scale_tril=self._tril),
            batch_shape=[2],
            event_shape=[3],
            validate_args=True)
  
    def testNonScalarBatchNonScalarEvent(self):
      # Can't override event_shape and/or batch_shape for non_scalar batch,
      # non-scalar event.
      with self.assertRaisesRegexp(ValueError, "base distribution not scalar"):
        self._cls()(
            distribution=tfd.MultivariateNormalDiag(
                loc=[[0.]], scale_diag=[[1.]]),
            bijector=tfb.Affine(shift=self._shift, scale_tril=self._tril),
            batch_shape=[2],
            event_shape=[3],
            validate_args=True)
  
    def testMatrixEvent(self):
      batch_shape = [2]
      event_shape = [2, 3, 3]
      batch_shape_pl = tf1.placeholder_with_default(
          input=np.int32(batch_shape), shape=None, name="dynamic_batch_shape")
      event_shape_pl = tf1.placeholder_with_default(
          input=np.int32(event_shape), shape=None, name="dynamic_event_shape")
  
      scale = 2.
      loc = 0.
      fake_mvn_dynamic = self._cls()(
          distribution=tfd.Normal(loc=loc, scale=scale),
          bijector=DummyMatrixTransform(),
          batch_shape=batch_shape_pl,
          event_shape=event_shape_pl,
          validate_args=True)
  
      fake_mvn_static = self._cls()(
          distribution=tfd.Normal(loc=loc, scale=scale),
          bijector=DummyMatrixTransform(),
          batch_shape=batch_shape,
          event_shape=event_shape,
          validate_args=True)
  
      def actual_mvn_log_prob(x):
        # This distribution is the normal PDF, reduced over the
        # last 3 dimensions + a jacobian term which corresponds
        # to the determinant of x.
        return (np.sum(stats.norm(loc, scale).logpdf(x), axis=(-1, -2, -3)) +
                np.sum(np.linalg.det(x), axis=-1))
  
      self.assertAllEqual([2, 3, 3], fake_mvn_static.event_shape)
      self.assertAllEqual([2], fake_mvn_static.batch_shape)
  
      if not tf.executing_eagerly():
        self.assertAllEqual(tf.TensorShape(None), fake_mvn_dynamic.event_shape)
        self.assertAllEqual(tf.TensorShape(None), fake_mvn_dynamic.batch_shape)
  
      num_samples = 5e3
      for fake_mvn in [fake_mvn_static, fake_mvn_dynamic]:
        # Ensure sample works by checking first, second moments.
        y = fake_mvn.sample(int(num_samples), seed=tfp_test_util.test_seed())
        x = y[0:5, ...]
        [
            x_,
            fake_event_shape_,
            fake_batch_shape_,
            fake_log_prob_,
            fake_prob_,
        ] = self.evaluate([
            x,
            fake_mvn.event_shape_tensor(),
            fake_mvn.batch_shape_tensor(),
            fake_mvn.log_prob(x),
            fake_mvn.prob(x),
        ])
  
        # Ensure all other functions work as intended.
        self.assertAllEqual([5, 2, 2, 3, 3], x_.shape)
        self.assertAllEqual([2, 3, 3], fake_event_shape_)
        self.assertAllEqual([2], fake_batch_shape_)
        self.assertAllClose(
            actual_mvn_log_prob(x_), fake_log_prob_, atol=0., rtol=1e-6)
        # With this many dimensions and samples, the direct space probability
        # may underflow.
        self.assertAllClose(
            np.exp(actual_mvn_log_prob(x_)), fake_prob_, atol=1e-12, rtol=1e-5)
  
    def testEmptyEvent(self):
      # Verify that zero-dimensional multivariate Normal distributions still
      # return reasonable shapes and a log-prob of 0.0.
      event_shape = [0]
      for batch_shape in ([2], []):
        for shapes_are_dynamic in (True, False):
          loc = tf.zeros(batch_shape + event_shape)
          scale_diag = tf.ones(batch_shape + event_shape)
          if shapes_are_dynamic:
            loc = tf1.placeholder_with_default(
                loc, shape=None, name="dynamic_loc")
            scale_diag = tf1.placeholder_with_default(
                scale_diag, shape=None, name="dynamic_scale_diag")
  
          mvn = tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale_diag)
  
          self.assertAllEqual(self.evaluate(mvn.event_shape_tensor()),
                              event_shape)
          self.assertAllEqual(self.evaluate(mvn.batch_shape_tensor()),
                              batch_shape)
          if not shapes_are_dynamic:
            self.assertAllEqual(
                tensorshape_util.as_list(mvn.event_shape), event_shape)
            self.assertAllEqual(
                tensorshape_util.as_list(mvn.batch_shape), batch_shape)
  
          for sample_shape in ([3], []):
            sample_ = self.evaluate(mvn.sample(sample_shape))
            self.assertAllEqual(sample_.shape,
                                sample_shape + batch_shape + event_shape)
            self.assertAllEqual(
                self.evaluate(mvn.log_prob(sample_)),
                np.zeros(sample_shape + batch_shape))
  
  
  if __name__ == "__main__":
    tf.test.main()
  