Internal change

PiperOrigin-RevId: 451046918

Author: pravnar

tensorflow_probability/python/bijectors/moyal_cdf.py

@@ -109,14 +109,15 @@ def _forward(self, x):
 
   def _inverse(self, y):
     with tf.control_dependencies(self._maybe_assert_valid_y(y)):
+      np_dtype = dtype_util.as_numpy_dtype(y.dtype)
       return (self.loc - self.scale *
-              (np.log(2.) + 2. * tf.math.log(tfp_math.erfcinv(y))))
+              (np.log(np_dtype(2.)) + 2. * tf.math.log(tfp_math.erfcinv(y))))
 
   def _inverse_log_det_jacobian(self, y):
     with tf.control_dependencies(self._maybe_assert_valid_y(y)):
-      return (tf.math.square(tfp_math.erfcinv(y)) +
-              tf.math.log(self.scale) + 0.5 * np.log(np.pi) -
-              tf.math.log(tfp_math.erfcinv(y)))
+      np_dtype = dtype_util.as_numpy_dtype(y.dtype)
+      return (tf.math.square(tfp_math.erfcinv(y)) + tf.math.log(self.scale) +
+              0.5 * np_dtype(np.log(np.pi)) - tf.math.log(tfp_math.erfcinv(y)))
 
   def _forward_log_det_jacobian(self, x):
     scale = tf.convert_to_tensor(self.scale)

tensorflow_probability/python/distributions/transformed_distribution.py

@@ -354,6 +354,13 @@ def _sample_and_log_prob(self, sample_shape, seed, **kwargs):
                tf.cast(fldj, base_distribution_log_prob.dtype))
 
   def _log_prob(self, y, **kwargs):
+    if self.bijector._is_injective:  # pylint: disable=protected-access
+      log_prob, _ = self.experimental_local_measure(
+          y, backward_compat=True, **kwargs)
+      return log_prob
+
+    # TODO(b/197680518): Support base measure handling for non-injective
+    # bijectors.
     distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)
 
     # For caching to work, it is imperative that the bijector is the first to
@@ -366,9 +373,6 @@ def _log_prob(self, y, **kwargs):
 
     ildj = self.bijector.inverse_log_det_jacobian(
         y, event_ndims=event_ndims, **bijector_kwargs)
-    if self.bijector._is_injective:  # pylint: disable=protected-access
-      base_log_prob = self.distribution.log_prob(x, **distribution_kwargs)
-      return base_log_prob + tf.cast(ildj, base_log_prob.dtype)
 
     # Compute log_prob on each element of the inverse image.
     lp_on_fibers = []
@@ -596,6 +600,32 @@ def _default_event_space_bijector(self):
         self.distribution.experimental_default_event_space_bijector())
   # pylint: enable=not-callable
 
+  def experimental_local_measure(self, y, backward_compat=False, **kwargs):
+    distribution_kwargs, bijector_kwargs = self._kwargs_split_fn(kwargs)
+
+    # For caching to work, it is imperative that the bijector is the first to
+    # modify the input.
+    x = self.bijector.inverse(y, **bijector_kwargs)
+    event_ndims = self.bijector.inverse_event_ndims(
+        tf.nest.map_structure(ps.rank_from_shape, self._event_shape_tensor(),
+                              self.event_shape), **bijector_kwargs)
+
+    if self.bijector._is_injective:  # pylint: disable=protected-access
+      local_measure_fn = self.distribution.experimental_local_measure
+      density_corr_fn = self.bijector.experimental_compute_density_correction
+      base_log_prob, tangent_space = local_measure_fn(
+          x, backward_compat=backward_compat, **distribution_kwargs)
+      correction, new_tangent_space = density_corr_fn(
+          x,
+          tangent_space,
+          backward_compat=backward_compat,
+          event_ndims=event_ndims,
+          **bijector_kwargs)
+      log_prob = base_log_prob - tf.cast(correction, base_log_prob.dtype)
+      return log_prob, new_tangent_space
+    else:
+      raise NotImplementedError
+
 
 class TransformedDistribution(
     _TransformedDistribution, distribution_lib.AutoCompositeTensorDistribution):
@@ -671,4 +701,3 @@ def _transformed_log_prob_ratio(p, x, q, y, name=None):
     ildj_ratio = ldj_ratio.inverse_log_det_jacobian_ratio(
         p.bijector, x, q.bijector, y, event_ndims)
     return base_log_prob_ratio + tf.cast(ildj_ratio, base_log_prob_ratio.dtype)
-

tensorflow_probability/python/experimental/bijectors/distribution_bijectors_test.py

@@ -42,12 +42,12 @@
     'LambertWNormal',  # CDF gradient incorrect at 0.
     'SigmoidBeta',  # inverse CDF numerical precision issues for large x
     'StudentT',  # CDF gradient incorrect at 0 (and unstable near zero).
-    )
+)
 
 if JAX_MODE:
   PRECONDITIONING_FAILS_DISTS = (
       'VonMises',  # Abstract eval for 'von_mises_cdf_jvp' not implemented.
-      ) + PRECONDITIONING_FAILS_DISTS
+  ) + PRECONDITIONING_FAILS_DISTS
 
 
 def _constrained_zeros_fn(shape, dtype, constraint_fn):
@@ -60,15 +60,18 @@ class DistributionBijectorsTest(test_util.TestCase):
 
   def assertDistributionIsApproximatelyStandardNormal(self,
                                                       dist,
+                                                      rtol=1e-6,
                                                       logprob_atol=1e-2,
                                                       grad_atol=1e-2):
     """Verifies that dist's lps and gradients match those of Normal(0., 1.)."""
     batch_shape = dist.batch_shape_tensor()
+
     def make_reference_values(event_shape):
       dist_shape = ps.concat([batch_shape, event_shape], axis=0)
       x = tf.reshape([-4., -2., 0., 2., 4.],
                      ps.concat([[5], ps.ones_like(dist_shape)], axis=0))
       return tf.broadcast_to(x, ps.concat([[5], dist_shape], axis=0))
+
     flat_event_shape = tf.nest.flatten(dist.event_shape_tensor())
     zs = [make_reference_values(s) for s in flat_event_shape]
     lp_dist, grad_dist = tfp.math.value_and_gradient(
@@ -83,11 +86,14 @@ def reference_value_and_gradient(z, event_shape):
     reference_vals_and_grads = [
         reference_value_and_gradient(z, event_shape)
         for (z, event_shape) in zip(zs, flat_event_shape)]
+
     lps_reference = [lp for lp, grad in reference_vals_and_grads]
-    self.assertAllClose(sum(lps_reference), lp_dist, atol=logprob_atol)
+    self.assertAllClose(
+        sum(lps_reference), lp_dist, rtol=rtol, atol=logprob_atol)
 
     grads_reference = [grad for lp, grad in reference_vals_and_grads]
-    self.assertAllCloseNested(grads_reference, grad_dist, atol=grad_atol)
+    self.assertAllCloseNested(
+        grads_reference, grad_dist, rtol=rtol, atol=grad_atol)
 
   @parameterized.named_parameters(
       {'testcase_name': dname, 'dist_name': dname}
@@ -101,10 +107,11 @@ def test_all_distributions_either_work_or_raise_error(self, dist_name, data):
     if dist_name in PRECONDITIONING_FAILS_DISTS:
       self.skipTest('Known failure.')
 
-    dist = data.draw(dhps.base_distributions(
-        dist_name=dist_name,
-        enable_vars=False,
-        param_strategy_fn=_constrained_zeros_fn))
+    dist = data.draw(
+        dhps.base_distributions(
+            dist_name=dist_name,
+            enable_vars=False,
+            param_strategy_fn=_constrained_zeros_fn))
     try:
       b = tfp.experimental.bijectors.make_distribution_bijector(dist)
     except NotImplementedError:
@@ -114,22 +121,20 @@ def test_all_distributions_either_work_or_raise_error(self, dist_name, data):
 
   @test_util.numpy_disable_gradient_test
   def test_multivariate_normal(self):
-    d = tfd.MultivariateNormalFullCovariance(loc=[4., 8.],
-                                             covariance_matrix=[[11., 0.099],
-                                                                [0.099, 0.1]])
+    d = tfd.MultivariateNormalFullCovariance(
+        loc=[4., 8.], covariance_matrix=[[11., 0.099], [0.099, 0.1]])
     b = tfp.experimental.bijectors.make_distribution_bijector(d)
-    self.assertDistributionIsApproximatelyStandardNormal(
-        tfb.Invert(b)(d))
+    self.assertDistributionIsApproximatelyStandardNormal(tfb.Invert(b)(d))
 
   @test_util.numpy_disable_gradient_test
   def test_markov_chain(self):
     d = tfd.MarkovChain(
         initial_state_prior=tfd.Uniform(low=0., high=1.),
         transition_fn=lambda _, x: tfd.Uniform(low=0., high=tf.nn.softplus(x)),
-        num_steps=10)
+        num_steps=3)
     b = tfp.experimental.bijectors.make_distribution_bijector(d)
     self.assertDistributionIsApproximatelyStandardNormal(
-        tfb.Invert(b)(d))
+        tfb.Invert(b)(d), rtol=1e-4)
 
   @test_util.numpy_disable_gradient_test
   def test_markov_chain_joint(self):
@@ -145,21 +150,22 @@ def test_markov_chain_joint(self):
         num_steps=10)
     b = tfp.experimental.bijectors.make_distribution_bijector(d)
     self.assertDistributionIsApproximatelyStandardNormal(
-        tfb.Invert(b)(d))
+        tfb.Invert(b)(d), rtol=1e-4)
 
   @test_util.numpy_disable_gradient_test
   def test_nested_joint_distribution(self):
 
     def model():
       x = yield tfd.Normal(loc=-2., scale=1.)
       yield tfd.JointDistributionSequentialAutoBatched([
-          tfd.Uniform(low=1. + tf.exp(x),
-                      high=1 + tf.exp(x) + tf.nn.softplus(x)),
+          tfd.Uniform(low=1. - tf.exp(x),
+                      high=2. + tf.exp(x) + tf.nn.softplus(x)),
           lambda v: tfd.Exponential(v)])  # pylint: disable=unnecessary-lambda
+
     dist = tfd.JointDistributionCoroutineAutoBatched(model)
     b = tfp.experimental.bijectors.make_distribution_bijector(dist)
     self.assertDistributionIsApproximatelyStandardNormal(
-        tfb.Invert(b)(dist))
+        tfb.Invert(b)(dist), rtol=1e-4)
 
   @test_util.numpy_disable_gradient_test
   @test_util.jax_disable_test_missing_functionality(
@@ -171,6 +177,7 @@ def model_with_funnel():
       z = yield tfd.Normal(loc=-1., scale=2., name='z')
       x = yield tfd.Normal(loc=[0.], scale=tf.exp(z), name='x')
       yield tfd.Poisson(log_rate=x, name='y')
+
     pinned_model = model_with_funnel.experimental_pin(y=[1])
     surrogate_posterior = tfp.experimental.vi.build_asvi_surrogate_posterior(
         pinned_model)
@@ -191,15 +198,16 @@ def do_sample():
           kernel=tfp.mcmc.DualAveragingStepSizeAdaptation(
               tfp.mcmc.TransformedTransitionKernel(
                   tfp.mcmc.NoUTurnSampler(
-                      pinned_model.unnormalized_log_prob,
-                      step_size=0.1),
+                      pinned_model.unnormalized_log_prob, step_size=0.1),
                   bijector=bijector),
               num_adaptation_steps=5),
           current_state=surrogate_posterior.sample(),
           num_burnin_steps=5,
           trace_fn=lambda _0, _1: [],
           num_results=10)
+
     do_sample()
 
+
 if __name__ == '__main__':
   test_util.main()

tensorflow_probability/python/experimental/util/jit_public_methods.py

@@ -36,6 +36,7 @@
     'dtype',
     'kl_divergence',  # Wrapping applied explicitly in `_traced_kl_divergence`.
     'experimental_default_event_space_bijector',
+    'experimental_local_measure',
     # tfb.Bijector
     # TODO(davmre): Test wrapping bijectors.
     'forward_event_shape',
@@ -45,7 +46,8 @@
     'forward_dtype',
     'inverse_dtype',
     'forward_event_ndims',
-    'inverse_event_ndims'
+    'inverse_event_ndims',
+    'experimental_compute_density_correction',
 )
 
 if NUMPY_MODE:

tensorflow_probability/python/sts/forecast_test.py

@@ -181,9 +181,8 @@ def _run():
 
   @test_util.jax_disable_test_missing_functionality('fit_with_hmc')
   def test_forecast_from_hmc(self):
-    if not (tf1.control_flow_v2_enabled() or self.use_static_shape):
-      self.skipTest('test_forecast_from_hmc does not currently work with TF1 '
-                    'and dynamic shapes')
+    if not tf1.control_flow_v2_enabled():
+      self.skipTest('test_forecast_from_hmc does not currently work with TF1')
 
     # test that we can directly plug in the output of an HMC chain as
     # the input to `forecast`, as done in the example, with no `sess.run` call.