impute_missing_values is not working

I just run the example from the official website:

I have encountered one warning and one error:

# Build model using observed time series to set heuristic priors.
linear_trend_model = tfp.sts.LocalLinearTrend(
  observed_time_series=observed_time_series)
model = tfp.sts.Sum([linear_trend_model],
                    observed_time_series=observed_time_series)

# Fit model to data
parameter_samples, _ = tfp.sts.fit_with_hmc(model, observed_time_series)

WARNING:tensorflow:From D:\xxx\xxx\Anaconda3\envs\pdm\lib\site-packages\tensorflow_probability\python\sts\fitting.py:561: SeedStream.init (from tensorflow_probability.python.util.seed_stream) is deprecated and will be removed after 2019-10-01.
Instructions for updating:
SeedStream has moved to tfp.util.SeedStream.
WARNING:tensorflow:From D:\xxx\xxx\Anaconda3\envs\pdm\lib\site-packages\tensorflow_core\python\ops\linalg\linear_operator_diag.py:166: calling LinearOperator.init (from tensorflow.python.ops.linalg.linear_operator) with graph_parents is deprecated and will be removed in a future version.
Instructions for updating:
Do not pass graph_parents. They will no longer be used.
WARNING:tensorflow:From D:\xxx\xxx\Anaconda3\envs\pdm\lib\site-packages\tensorflow_core\python\ops\linalg\linear_operator_block_diag.py:199: LinearOperator.graph_parents (from tensorflow.python.ops.linalg.linear_operator) is deprecated and will be removed in a future version.
Instructions for updating:
Do not call graph_parents.

# Impute missing values
imputed_series_distribution = tfp.sts.impute_missing_values(
  model, observed_time_series)
print('imputed means and stddevs: ',
      imputed_series_distribution.mean(),
      imputed_series_distribution.stddev())

TypeError Traceback (most recent call last)
in
1 # Impute missing values
2 imputed_series_distribution = tfp.sts.impute_missing_values(
----> 3 model, observed_time_series)
4 print('imputed means and stddevs: ',
5 imputed_series_distribution.mean(),

TypeError: impute_missing_values() missing 1 required positional argument: 'parameter_samples'

Then I tried to fill the missing argument:

# Impute missing values
imputed_series_distribution = tfp.sts.impute_missing_values(
  model, observed_time_series, parameter_samples=None)
print('imputed means and stddevs: ',
      imputed_series_distribution.mean(),
      imputed_series_distribution.stddev())

TypeError Traceback (most recent call last)
in
1 # Impute missing values
2 imputed_series_distribution = tfp.sts.impute_missing_values(
----> 3 model, observed_time_series, parameter_samples=None)
4 print('imputed means and stddevs: ',
5 imputed_series_distribution.mean(),

D:\xxx\xxx\Anaconda3\envs\pdm\lib\site-packages\tensorflow_probability\python\sts\forecast.py in impute_missing_values(model, observed_time_series, parameter_samples, include_observation_noise)
449 tf.shape(input=observed_time_series))[-2]
450 lgssm = model.make_state_space_model(
--> 451 num_timesteps=num_timesteps, param_vals=parameter_samples)
452 posterior_means, posterior_covs = lgssm.posterior_marginals(
453 observed_time_series, mask=mask)

D:\xxx\xxx\Anaconda3\envs\pdm\lib\site-packages\tensorflow_probability\python\sts\structural_time_series.py in make_state_space_model(self, num_timesteps, param_vals, initial_state_prior, initial_step)
156 return self._make_state_space_model(
157 num_timesteps=num_timesteps,
--> 158 param_map=self._canonicalize_param_vals_as_map(param_vals),
159 initial_state_prior=initial_state_prior,
160 initial_step=initial_step)

D:\xxx\xxx\Anaconda3\envs\pdm\lib\site-packages\tensorflow_probability\python\sts\structural_time_series.py in _canonicalize_param_vals_as_map(self, param_vals)
129 param_map = param_vals
130 else:
--> 131 param_map = {p.name: v for (p, v) in zip(self.parameters, param_vals)}
132
133 return param_map

TypeError: zip argument #2 must support iteration

Hope the example code from the official instruction page can be fixed.

Many thanks

Start with the linked TensorFlow Probability API example for tfp.sts.impute_missing_values and compare it with the function signature shown in the traceback. Reproduce the notebook steps around fit_with_hmc and parameter_samples, then update the example so it runs successfully and verify that the documented imputed means and standard deviations are produced.