Add another structured state test, then make it pass with a small change to PHMC.

PiperOrigin-RevId: 346322156

Author: brianwa84

tensorflow_probability/python/experimental/mcmc/preconditioned_hmc.py

@@ -280,7 +280,7 @@ def one_step(self, current_state, previous_kernel_results, seed=None):
           state_gradients_are_stopped=self.state_gradients_are_stopped)
 
       seed = samplers.sanitize_seed(seed)
-      current_momentum_parts = momentum_distribution.sample(seed=seed)
+      current_momentum_parts = list(momentum_distribution.sample(seed=seed))
       momentum_log_prob = getattr(momentum_distribution,
                                   '_log_prob_unnormalized',
                                   momentum_distribution.log_prob)

tensorflow_probability/python/mcmc/internal/util.py

@@ -144,7 +144,8 @@ def make_name(super_name, default_super_name, sub_name):
 def _choose_base_case(is_accepted,
                       proposed,
                       current,
-                      name=None):
+                      name=None,
+                      addr=None,):
   """Helper to `choose` which expand_dims `is_accepted` and applies tf.where."""
   def _where(proposed, current):
     """Wraps `tf.where`."""
@@ -162,30 +163,38 @@ def _where(proposed, current):
   with tf.name_scope(name or 'choose'):
     if not is_list_like(proposed):
       return _where(proposed, current)
-    return [(choose(is_accepted, p, c, name=name) if is_namedtuple_like(p)
-             else _where(p, c))
-            for p, c in zip(proposed, current)]
+    return tf.nest.pack_sequence_as(
+        current,
+        [(_choose_recursive(is_accepted, p, c, name=name, addr=f'{addr}[i]')
+          if is_namedtuple_like(p) else
+          _where(p, c)) for i, (p, c) in enumerate(zip(proposed, current))])
 
 
-def choose(is_accepted, proposed, current, name=None):
-  """Helper which expand_dims `is_accepted` then applies tf.where."""
+def _choose_recursive(is_accepted, proposed, current, name=None, addr='<root>'):
+  """Recursion helper which also reports the address of any failures."""
   with tf.name_scope(name or 'choose'):
     if not is_namedtuple_like(proposed):
-      return _choose_base_case(is_accepted, proposed, current, name=name)
+      return _choose_base_case(is_accepted, proposed, current, name=name,
+                               addr=addr)
     if not isinstance(proposed, type(current)):
-      raise TypeError('Type of `proposed` ({}) must be identical to '
-                      'type of `current` ({})'.format(
-                          type(proposed).__name__,
-                          type(current).__name__))
+      raise TypeError(
+          f'Type of `proposed` ({type(proposed).__name__}) must be identical '
+          f'to type of `current` ({type(current).__name__}). (At "{addr}".)')
     items = {}
     for fn in proposed._fields:
-      items[fn] = choose(is_accepted,
-                         getattr(proposed, fn),
-                         getattr(current, fn),
-                         name=name)
+      items[fn] = _choose_recursive(is_accepted,
+                                    getattr(proposed, fn),
+                                    getattr(current, fn),
+                                    name=name,
+                                    addr=f'{addr}/{fn}')
     return type(proposed)(**items)
 
 
+def choose(is_accepted, proposed, current, name=None):
+  """Helper which expand_dims `is_accepted` then applies tf.where."""
+  return _choose_recursive(is_accepted, proposed, current, name=name)
+
+
 def strip_seeds(obj):
   if not is_namedtuple_like(obj):
     return obj

tensorflow_probability/python/mcmc/sample_test.py

@@ -31,6 +31,7 @@
 from tensorflow_probability.python.internal import tensorshape_util
 from tensorflow_probability.python.internal import test_util
 
+tfb = tfp.bijectors
 tfd = tfp.distributions
 
 
@@ -404,6 +405,40 @@ def sample():
           seed=seed_stream())
     self.evaluate(sample())
 
+  @test_util.jax_disable_test_missing_functionality('PHMC b/175107050')
+  @test_util.numpy_disable_gradient_test('HMC')
+  def testStructuredState2(self):
+    @tfd.JointDistributionCoroutineAutoBatched
+    def model():
+      mu = yield tfd.Sample(tfd.Normal(0, 1), [65], name='mu')
+      sigma = yield tfd.Sample(tfd.Exponential(1.), [65], name='sigma')
+      beta = yield tfd.Sample(
+          tfd.Normal(loc=tf.gather(mu, tf.range(436) % 65, axis=-1),
+                     scale=tf.gather(sigma, tf.range(436) % 65, axis=-1)),
+          4, name='beta')
+      _ = yield tfd.Multinomial(total_count=100.,
+                                logits=tfb.Pad([[0, 1]])(beta),
+                                name='y')
+
+    stream = test_util.test_seed_stream()
+    pinned = model.experimental_pin(y=model.sample(seed=stream()).y)
+    struct = pinned.dtype
+    stddevs = struct._make([
+        tf.fill([65], .1), tf.fill([65], 1.), tf.fill([436, 4], 10.)])
+    momentum_dist = tfd.JointDistributionNamedAutoBatched(
+        struct._make(tfd.Normal(0, 1 / std) for std in stddevs))
+    kernel = tfp.experimental.mcmc.PreconditionedHamiltonianMonteCarlo(
+        pinned.unnormalized_log_prob,
+        step_size=.1, num_leapfrog_steps=10,
+        momentum_distribution=momentum_dist)
+    bijector = pinned.experimental_default_event_space_bijector()
+    kernel = tfp.mcmc.TransformedTransitionKernel(kernel, bijector)
+    state = bijector(struct._make(
+        tfd.Uniform(-2., 2.).sample(shp)
+        for shp in bijector.inverse_event_shape(pinned.event_shape)))
+    self.evaluate(tfp.mcmc.sample_chain(
+        3, current_state=state, kernel=kernel, seed=stream()))
+
 
 if __name__ == '__main__':
   tf.test.main()