Add use_markov_chain flag to the inference gym's Brownian motion model.

PiperOrigin-RevId: 385001418

Author: davmre

spinoffs/inference_gym/inference_gym/targets/BUILD

@@ -128,8 +128,10 @@ py_test(
     srcs_version = "PY3",
     deps = [
         ":brownian_motion",
+        # absl/testing:parameterized dep,
         # numpy dep,
         # tensorflow dep,
+        # tensorflow_probability/python/internal:test_util dep,
         "//inference_gym/internal:test_util",
     ],
 )

spinoffs/inference_gym/inference_gym/targets/brownian_motion.py

@@ -38,7 +38,17 @@
 Root = tfd.JointDistributionCoroutine.Root
 
 
-def brownian_motion_prior_fn(num_timesteps, innovation_noise_scale):
+def brownian_motion_as_markov_chain(num_timesteps, innovation_noise_scale):
+  return tfd.MarkovChain(
+      initial_state_prior=tfd.Normal(loc=0., scale=innovation_noise_scale),
+      transition_fn=lambda _, x_t: tfd.Normal(  # pylint: disable=g-long-lambda
+          loc=x_t, scale=innovation_noise_scale),
+      num_steps=num_timesteps,
+      name='locs')
+
+
+def brownian_motion_prior_fn(num_timesteps,
+                             innovation_noise_scale):
   """Generative process for the Brownian Motion model."""
   prior_loc = 0.
   new = yield Root(tfd.Normal(loc=prior_loc,
@@ -50,25 +60,34 @@ def brownian_motion_prior_fn(num_timesteps, innovation_noise_scale):
                            name='x_{}'.format(t))
 
 
-def brownian_motion_unknown_scales_prior_fn(num_timesteps):
+def brownian_motion_unknown_scales_prior_fn(num_timesteps, use_markov_chain):
   """Generative process for the Brownian Motion model with unknown scales."""
   innovation_noise_scale = yield Root(tfd.LogNormal(
       0., 2., name='innovation_noise_scale'))
   _ = yield Root(tfd.LogNormal(0., 2., name='observation_noise_scale'))
-  yield from brownian_motion_prior_fn(
-      num_timesteps,
-      innovation_noise_scale=innovation_noise_scale)
+  if use_markov_chain:
+    yield brownian_motion_as_markov_chain(
+        num_timesteps=num_timesteps,
+        innovation_noise_scale=innovation_noise_scale)
+  else:
+    yield from brownian_motion_prior_fn(
+        num_timesteps,
+        innovation_noise_scale=innovation_noise_scale)
 
 
 def brownian_motion_log_likelihood_fn(values,
                                       observed_locs,
+                                      use_markov_chain,
                                       observation_noise_scale=None):
   """Likelihood of observed data under the Brownian Motion model."""
   if observation_noise_scale is None:
-    (_, observation_noise_scale), values = values[:2], values[2:]
+    (_, observation_noise_scale) = values[:2]
+    latents = values[2] if use_markov_chain else tf.stack(values[2:], axis=-1)
+  else:
+    latents = values if use_markov_chain else tf.stack(values, axis=-1)
+
   observation_noise_scale = tf.convert_to_tensor(
       observation_noise_scale, name='observation_noise_scale')
-  latents = tf.stack(values, axis=-1)
   is_observed = ~tf.math.is_nan(observed_locs)
   lps = tfd.Normal(
       loc=latents, scale=observation_noise_scale[..., tf.newaxis]).log_prob(
@@ -98,6 +117,7 @@ def __init__(self,
                observed_locs,
                innovation_noise_scale,
                observation_noise_scale,
+               use_markov_chain=False,
                name='brownian_motion',
                pretty_name='Brownian Motion'):
     """Construct the Brownian Motion model.
@@ -107,35 +127,52 @@ def __init__(self,
         unobserved.
       innovation_noise_scale: Python `float`.
       observation_noise_scale: Python `float`.
+      use_markov_chain: Python `bool` indicating whether to use the
+        `MarkovChain` distribution in place of separate random variables for
+        each time step. The default of `False` is for backwards compatibility;
+        setting this to `True` should significantly improve performance.
       name: Python `str` name prefixed to Ops created by this class.
       pretty_name: A Python `str`. The pretty name of this model.
     """
     with tf.name_scope(name):
       num_timesteps = observed_locs.shape[0]
-      self._prior_dist = tfd.JointDistributionCoroutine(
-          functools.partial(
-              brownian_motion_prior_fn,
-              num_timesteps=num_timesteps,
-              innovation_noise_scale=innovation_noise_scale))
+      if use_markov_chain:
+        self._prior_dist = brownian_motion_as_markov_chain(
+            num_timesteps=num_timesteps,
+            innovation_noise_scale=innovation_noise_scale)
+      else:
+        self._prior_dist = tfd.JointDistributionCoroutine(
+            functools.partial(
+                brownian_motion_prior_fn,
+                num_timesteps=num_timesteps,
+                innovation_noise_scale=innovation_noise_scale))
 
       self._log_likelihood_fn = functools.partial(
           brownian_motion_log_likelihood_fn,
           observation_noise_scale=observation_noise_scale,
-          observed_locs=observed_locs)
+          observed_locs=observed_locs,
+          use_markov_chain=use_markov_chain)
 
       def _ext_identity(params):
         return tf.stack(params, axis=-1)
 
+      def _ext_identity_markov_chain(params):
+        return params
+
       sample_transformations = {
           'identity':
               model.Model.SampleTransformation(
-                  fn=_ext_identity,
+                  fn=(_ext_identity_markov_chain
+                      if use_markov_chain else _ext_identity),
                   pretty_name='Identity',
               )
       }
 
-    event_space_bijector = type(
-        self._prior_dist.dtype)(*([tfb.Identity()] * num_timesteps))
+    if use_markov_chain:
+      event_space_bijector = tfb.Identity()
+    else:
+      event_space_bijector = type(
+          self._prior_dist.dtype)(*([tfb.Identity()] * num_timesteps))
     super(BrownianMotion, self).__init__(
         default_event_space_bijector=event_space_bijector,
         event_shape=self._prior_dist.event_shape,
@@ -157,11 +194,12 @@ class BrownianMotionMissingMiddleObservations(BrownianMotion):
 
   GROUND_TRUTH_MODULE = brownian_motion_missing_middle_observations
 
-  def __init__(self):
+  def __init__(self, use_markov_chain=False):
     dataset = data.brownian_motion_missing_middle_observations()
     super(BrownianMotionMissingMiddleObservations, self).__init__(
         name='brownian_motion_missing_middle_observations',
         pretty_name='Brownian Motion Missing Middle Observations',
+        use_markov_chain=use_markov_chain,
         **dataset)
 
 
@@ -188,13 +226,19 @@ class BrownianMotionUnknownScales(bayesian_model.BayesianModel):
 
   def __init__(self,
                observed_locs,
+               use_markov_chain=False,
                name='brownian_motion_unknown_scales',
                pretty_name='Brownian Motion with Unknown Scales'):
     """Construct the Brownian Motion model with unknown scales.
 
     Args:
       observed_locs: Array of loc parameters with nan value if loc is
         unobserved.
+      use_markov_chain: Python `bool` indicating whether to use the
+        `MarkovChain` distribution in place of separate random variables for
+        each time step. The default of `False` is for backwards compatibility;
+        setting this to `True` should significantly improve performance.
+        Default value: `False`.
       name: Python `str` name prefixed to Ops created by this class.
       pretty_name: A Python `str`. The pretty name of this model.
     """
@@ -203,16 +247,20 @@ def __init__(self,
       self._prior_dist = tfd.JointDistributionCoroutine(
           functools.partial(
               brownian_motion_unknown_scales_prior_fn,
+              use_markov_chain=use_markov_chain,
               num_timesteps=num_timesteps))
 
       self._log_likelihood_fn = functools.partial(
           brownian_motion_log_likelihood_fn,
+          use_markov_chain=use_markov_chain,
           observed_locs=observed_locs)
 
       def _ext_identity(params):
         return {'innovation_noise_scale': params[0],
                 'observation_noise_scale': params[1],
-                'locs': tf.stack(params[2:], axis=-1)}
+                'locs': (params[2]
+                         if use_markov_chain
+                         else tf.stack(params[2:], axis=-1))}
 
       sample_transformations = {
           'identity':
@@ -228,7 +276,8 @@ def _ext_identity(params):
         self._prior_dist.dtype)(*(
             [tfb.Softplus(),
              tfb.Softplus()
-             ] + [tfb.Identity()] * num_timesteps))
+             ] + [tfb.Identity()] * (
+                 1 if use_markov_chain else num_timesteps)))
     super(BrownianMotionUnknownScales, self).__init__(
         default_event_space_bijector=event_space_bijector,
         event_shape=self._prior_dist.event_shape,
@@ -252,11 +301,12 @@ class BrownianMotionUnknownScalesMissingMiddleObservations(
   GROUND_TRUTH_MODULE = (
       brownian_motion_unknown_scales_missing_middle_observations)
 
-  def __init__(self):
+  def __init__(self, use_markov_chain=False):
     dataset = data.brownian_motion_missing_middle_observations()
     del dataset['innovation_noise_scale']
     del dataset['observation_noise_scale']
     super(BrownianMotionUnknownScalesMissingMiddleObservations, self).__init__(
         name='brownian_motion_unknown_scales_missing_middle_observations',
         pretty_name='Brownian Motion with Unknown Scales',
+        use_markov_chain=use_markov_chain,
         **dataset)

spinoffs/inference_gym/inference_gym/targets/brownian_motion_test.py

@@ -19,6 +19,7 @@
 import numpy as np
 import tensorflow.compat.v2 as tf
 
+from tensorflow_probability.python.internal import test_util as tfp_test_util
 from inference_gym.internal import test_util
 from inference_gym.targets import brownian_motion
 
@@ -112,5 +113,36 @@ def testBrownianMotionUnknownScalesMissingMiddleObservationsHMC(self):
         step_size=0.03,
     )
 
+  def testBrownianMotionMarkovChainLogprobMatchesOriginal(self):
+    model = (
+        brownian_motion.BrownianMotionMissingMiddleObservations(
+            use_markov_chain=False))
+    markov_chain_model = (
+        brownian_motion.BrownianMotionMissingMiddleObservations(
+            use_markov_chain=True))
+
+    x = self.evaluate(model.prior_distribution().sample(
+        400, seed=tfp_test_util.test_seed()))
+    self.assertAllClose(model.unnormalized_log_prob(x),
+                        markov_chain_model.unnormalized_log_prob(
+                            tf.stack(x, axis=-1)),
+                        atol=1e-2)
+
+  def testBrownianMotionUnknownScalesMarkovChainLogprobMatchesOriginal(self):
+    model = (
+        brownian_motion.BrownianMotionUnknownScalesMissingMiddleObservations(
+            use_markov_chain=False))
+    markov_chain_model = (
+        brownian_motion.BrownianMotionUnknownScalesMissingMiddleObservations(
+            use_markov_chain=True))
+
+    x = self.evaluate(model.prior_distribution().sample(
+        400, seed=tfp_test_util.test_seed()))
+    self.assertAllClose(model.unnormalized_log_prob(x),
+                        markov_chain_model.unnormalized_log_prob(
+                            type(markov_chain_model.dtype)(
+                                x[0], x[1], tf.stack(x[2:], axis=-1))),
+                        atol=1e-2)
+
 if __name__ == '__main__':
   tf.test.main()