Fix failure in windowed_sampling_test.jax in OSS.

PiperOrigin-RevId: 453202511
tensorflow · Jun 6, 2022 · 8e72c11 · 8e72c11
1 parent f8107c1
commit 8e72c11
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/tensorflow_probability/python/experimental/mcmc/windowed_sampling.py b/tensorflow_probability/python/experimental/mcmc/windowed_sampling.py
@@ -269,8 +269,9 @@ def step_broadcast(step_size):
   shard_axis_names = pinned_model.experimental_shard_axis_names
   if any(tf.nest.flatten(shard_axis_names)):
     shard_axis_names = nest.flatten_up_to(
-        initial_transformed_position, pinned_model._model_flatten(  # pylint: disable=protected-access
-            shard_axis_names))
+        initial_transformed_position,
+        list(pinned_model._model_flatten(shard_axis_names)))  # pylint: disable=protected-access
+
   else:
     # No active shard axis names
     shard_axis_names = None