Fixed problem where the step size in NUTS and preconditioned NUTS could have a different structure than what the user input.

Specifically, when step size was a scalar, using `trace_fn`, or `return_final_kernel_results=True`, the step size would be a list. This did not happen in HMC (or preconditioned HMC), and this change makes the behavior consistent.

PiperOrigin-RevId: 374703573

Author: ColCarroll

tensorflow_probability/python/experimental/mcmc/pnuts_test.py

@@ -589,6 +589,15 @@ def trace_fn(_, pkr):
     self.assertAllClose(
         average_rhat, np.ones_like(average_rhat), atol=0.05, rtol=0.05)
 
+  def test_step_size_trace(self):
+    dist = tfd.Normal(0., 1.)
+    kernel = tfp.experimental.mcmc.PreconditionedNoUTurnSampler(
+        dist.log_prob, step_size=1.)
+    _, _, fkr = tfp.mcmc.sample_chain(10, 0., kernel=kernel,
+                                      return_final_kernel_results=True,
+                                      seed=test_util.test_seed())
+    self.assertAlmostEqual(1., self.evaluate(fkr.step_size))
+
 # Allowed type of preconditioning schemes to use.
 # See code for details.
 PRECONDITION_SCHEMES = frozenset([

tensorflow_probability/python/experimental/mcmc/preconditioned_nuts.py

@@ -270,8 +270,6 @@ def __init__(self,
 
       # Process all other arguments.
       self._target_log_prob_fn = target_log_prob_fn
-      if not tf.nest.is_nested(step_size):
-        step_size = [step_size]
       self._step_size = step_size
 
       self._parameters = dict(
@@ -409,12 +407,16 @@ def _copy(v):
           read_instruction=read_instruction
           )
 
+      step_size = _prepare_step_size(
+          previous_kernel_results.step_size,
+          current_target_log_prob.dtype,
+          len(current_state))
       _, _, _, new_step_metastate = tf.while_loop(
           cond=lambda iter_, seed, state, metastate: (  # pylint: disable=g-long-lambda
               (iter_ < self.max_tree_depth) &
               tf.reduce_any(metastate.continue_tree)),
           body=lambda iter_, seed, state, metastate: self._loop_tree_doubling(  # pylint: disable=g-long-lambda
-              previous_kernel_results.step_size,
+              step_size,
               previous_kernel_results.velocity_state_memory,
               current_step_meta_info,
               iter_,
@@ -466,21 +468,9 @@ def bootstrap_results(self, init_state):
                                                      name='current_state')
       current_target_log_prob, current_grads_log_prob = mcmc_util.maybe_call_fn_and_grads(
           self.target_log_prob_fn, state_parts)
-      # Padding the step_size so it is compatable with the states
-      step_size = self.step_size
-      if len(step_size) == 1:
-        step_size = step_size * len(init_state)
-      if len(step_size) != len(init_state):
-        raise ValueError('Expected either one step size or {} (size of '
-                         '`init_state`), but found {}'.format(
-                             len(init_state), len(step_size)))
-      step_size = tf.nest.map_structure(
-          lambda x: tf.convert_to_tensor(  # pylint: disable=g-long-lambda
-              x,
-              dtype=current_target_log_prob.dtype,
-              name='step_size'),
-          step_size)
-
+      # Confirm that the step size is compatible with the state parts.
+      _ = _prepare_step_size(
+          self.step_size, current_target_log_prob.dtype, len(init_state))
       momentum_distribution = self.momentum_distribution
       if momentum_distribution is None:
         momentum_distribution = pu.make_momentum_distribution(
@@ -508,7 +498,12 @@ def _init(shape_and_dtype):
           target_log_prob=current_target_log_prob,
           grads_target_log_prob=current_grads_log_prob,
           velocity_state_memory=velocity_state_memory,
-          step_size=step_size,
+          step_size=tf.nest.map_structure(
+              lambda x: tf.convert_to_tensor(  # pylint: disable=g-long-lambda
+                  x,
+                  dtype=current_target_log_prob.dtype,
+                  name='step_size'),
+              self.step_size),
           log_accept_ratio=tf.zeros_like(
               current_target_log_prob, name='log_accept_ratio'),
           leapfrogs_taken=tf.zeros_like(
@@ -1110,3 +1105,14 @@ def compute_hamiltonian(target_log_prob, momentum_parts, momentum_distribution):
 def get_kinetic_energy_fn(momentum_distribution):
   """Convert a momentum distribution to a kinetic energy function."""
   return lambda *args: -momentum_distribution.log_prob(*args)
+
+
+def _prepare_step_size(step_size, dtype, n_state_parts):
+  step_sizes, _ = mcmc_util.prepare_state_parts(
+      step_size, dtype=dtype, name='step_size')
+  if len(step_sizes) == 1:
+    step_sizes *= n_state_parts
+  if n_state_parts != len(step_sizes):
+    raise ValueError('There should be exactly one `step_size` or it should '
+                     'have same length as `current_state`.')
+  return step_sizes

tensorflow_probability/python/mcmc/nuts.py

@@ -262,8 +262,6 @@ def __init__(self,
 
       # Process all other arguments.
       self._target_log_prob_fn = target_log_prob_fn
-      if not tf.nest.is_nested(step_size):
-        step_size = [step_size]
       self._step_size = step_size
 
       self._parameters = dict(
@@ -400,12 +398,16 @@ def _copy(v):
           read_instruction=read_instruction
           )
 
+      step_size = _prepare_step_size(
+          previous_kernel_results.step_size,
+          current_target_log_prob.dtype,
+          len(current_state))
       _, _, _, new_step_metastate = tf.while_loop(
           cond=lambda iter_, seed, state, metastate: (  # pylint: disable=g-long-lambda
               (iter_ < self.max_tree_depth) &
               tf.reduce_any(metastate.continue_tree)),
           body=lambda iter_, seed, state, metastate: self._loop_tree_doubling(  # pylint: disable=g-long-lambda
-              previous_kernel_results.step_size,
+              step_size,
               previous_kernel_results.momentum_state_memory,
               current_step_meta_info,
               iter_,
@@ -472,26 +474,20 @@ def _init(shape_and_dtype):
       ] = leapfrog_impl.process_args(self.target_log_prob_fn, dummy_momentum,
                                      init_state)
 
-      # Padding the step_size so it is compatable with the states
-      step_size = self.step_size
-      if len(step_size) == 1:
-        step_size = step_size * len(init_state)
-      if len(step_size) != len(init_state):
-        raise ValueError('Expected either one step size or {} (size of '
-                         '`init_state`), but found {}'.format(
-                             len(init_state), len(step_size)))
-      step_size = tf.nest.map_structure(
-          lambda x: tf.convert_to_tensor(  # pylint: disable=g-long-lambda
-              x,
-              dtype=current_target_log_prob.dtype,
-              name='step_size'),
-          step_size)
+      # Confirm that the step size is compatible with the state parts.
+      _ = _prepare_step_size(
+          self.step_size, current_target_log_prob.dtype, len(init_state))
 
       return NUTSKernelResults(
           target_log_prob=current_target_log_prob,
           grads_target_log_prob=current_grads_log_prob,
           momentum_state_memory=momentum_state_memory,
-          step_size=step_size,
+          step_size=tf.nest.map_structure(
+              lambda x: tf.convert_to_tensor(  # pylint: disable=g-long-lambda
+                  x,
+                  dtype=current_target_log_prob.dtype,
+                  name='step_size'),
+              self.step_size),
           log_accept_ratio=tf.zeros_like(current_target_log_prob,
                                          name='log_accept_ratio'),
           leapfrogs_taken=tf.zeros_like(current_target_log_prob,
@@ -1080,6 +1076,17 @@ def generate_efficient_write_read_instruction(instruction_array):
   return write_instruction, np.asarray(read_instruction)
 
 
+def _prepare_step_size(step_size, dtype, n_state_parts):
+  step_sizes, _ = mcmc_util.prepare_state_parts(
+      step_size, dtype=dtype, name='step_size')
+  if len(step_sizes) == 1:
+    step_sizes *= n_state_parts
+  if n_state_parts != len(step_sizes):
+    raise ValueError('There should be exactly one `step_size` or it should '
+                     'have same length as `current_state`.')
+  return step_sizes
+
+
 def compute_hamiltonian(target_log_prob, momentum_parts,
                         shard_axis_names=None):
   """Compute the Hamiltonian of the current system."""

tensorflow_probability/python/mcmc/nuts_test.py

@@ -574,6 +574,14 @@ def trace_fn(_, pkr):
     self.assertAllClose(
         average_rhat, np.ones_like(average_rhat), atol=0.05, rtol=0.05)
 
+  def test_step_size_trace(self):
+    dist = tfd.Normal(0., 1.)
+    kernel = tfp.mcmc.NoUTurnSampler(dist.log_prob, step_size=1.)
+    _, _, fkr = tfp.mcmc.sample_chain(10, 0., kernel=kernel,
+                                      return_final_kernel_results=True,
+                                      seed=test_util.test_seed())
+    self.assertAlmostEqual(1., self.evaluate(fkr.step_size))
+
 
 @test_util.test_all_tf_execution_regimes
 class DistributedNutsTest(distribute_test_lib.DistributedTest):