Fix handling of scalar argument to init_step_size in windowed samplers.

This had been erroneously broadcast up to the shape of the state part, which meant there was no way to have a single step size shared between all state parts, other than relying on the default.

PiperOrigin-RevId: 377331772

Author: ColCarroll

tensorflow_probability/python/experimental/mcmc/windowed_sampling.py

@@ -255,9 +255,13 @@ def target_log_prob_fn(*args):
     return lp + ldj
 
   def step_broadcast(step_size):
-    return step_bijector(
-        nest_util.broadcast_structure(pinned_model.event_shape_tensor(),
-                                      step_size))
+    # Only apply the bijector to nested step sizes or non-scalar batches.
+    if tf.nest.is_nested(step_size):
+      return step_bijector(
+          nest_util.broadcast_structure(pinned_model.event_shape_tensor(),
+                                        step_size))
+    else:
+      return step_size
 
   return (target_log_prob_fn,
           initial_transformed_position,

tensorflow_probability/python/experimental/mcmc/windowed_sampling_test.py

@@ -428,8 +428,7 @@ def model_fn():
         num_adaptation_steps=100, init_step_size=tf.ones([10, 1]),
         seed=test_util.test_seed()))
     self.assertEqual((2, 64, 10, 3), states.x.shape)
-    self.assertLen(trace['step_size'], 1)
-    self.assertEqual((2, 10, 1), trace['step_size'][0].shape)
+    self.assertEqual((2, 10, 1), trace['step_size'].shape)
 
   def test_batch_of_problems_named(self):
     use_multinomial = tf.executing_eagerly()
@@ -457,8 +456,7 @@ def mk_y(x):
             init_step_size=tf.ones([10, 1]),
             seed=test_util.test_seed()))
     self.assertEqual((2, 64, 10, 3), states['x'].shape)
-    self.assertLen(trace['step_size'], 1)
-    self.assertEqual((2, 10, 1), trace['step_size'][0].shape)
+    self.assertEqual((2, 10, 1), trace['step_size'].shape)
 
 
 @test_util.test_graph_and_eager_modes
@@ -529,21 +527,20 @@ def test_supply_single_step_size(self):
     })
 
     init_step_size = 1.
-    _, actual_step_size = tfp.experimental.mcmc.windowed_adaptive_hmc(
-        1,
-        jd_model,
-        num_adaptation_steps=100,
-        n_chains=20,
-        init_step_size=init_step_size,
-        num_leapfrog_steps=5,
-        discard_tuning=False,
-        trace_fn=lambda *args: unnest.get_innermost(args[-1], 'step_size'),
-        seed=stream(),
-    )
-
-    actual_step = [j[0] for j in actual_step_size]
-    expected_step = [1., 1.]
-    self.assertAllCloseNested(expected_step, actual_step)
+    _, traced_step_size = self.evaluate(
+        tfp.experimental.mcmc.windowed_adaptive_hmc(
+            1,
+            jd_model,
+            num_adaptation_steps=100,
+            n_chains=20,
+            init_step_size=init_step_size,
+            num_leapfrog_steps=5,
+            discard_tuning=False,
+            trace_fn=lambda *args: unnest.get_innermost(args[-1], 'step_size'),
+            seed=stream()))
+
+    self.assertEqual((100 + 1,), traced_step_size.shape)
+    self.assertAllClose(1., traced_step_size[0])
 
   def test_sequential_step_size(self):
     stream = test_util.test_seed_stream()