Speed of tfp.sts.fit_with_hmc

Hi,

I've been building some structural time-series models with Tensorflow Probability over the past week.

I've begun to look into the impute_missing_variable method to smooth in my missing time-series values, but the sts.fit_with_hmc seems incredibly slow. Just running the example on https://www.tensorflow.org/probability/api_docs/python/tfp/sts/impute_missing_values take about 25 seconds.

import numpy as np
import  tensorflow_probability as tfp
import tensorflow as tf

time_series_with_nans = [-1., 1., np.nan, 2.4, np.nan, 5]
observed_time_series = tfp.sts.MaskedTimeSeries(
  time_series=time_series_with_nans,
    is_missing=tf.math.is_nan(time_series_with_nans))

# 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)

Is this expected behavior? I'm running this example on CPU and didn't see any performance improvements when I ran the same block of code on GPU.

This is using Tensorflow-Probability==0.19 and Tensorflow==2.11.0

Thanks!

Start by running the TensorFlow Probability impute_missing_values example and the tfp.sts.fit_with_hmc call on CPU and GPU as reported. Compare the observed runtime and inspect the fit_with_hmc entry point to determine whether the delay is expected or indicates a performance issue. Done means documenting the cause or providing a measurable improvement.