problems with tfp.sts.fit_with_hmc

I'm having a few issues with tfp.sts.fit_with_hmc so thought I would reach out, I am using a very recent tfp-nightly and tensorflow 2 alpha.

Firstly, the function signature has num_variational_steps=150 but the documentation has "Default value: 200".

Secondly, I've tried a simple model to try and understand things, but tfp.sts.fit_with_hmc is taking over 2 minutes for 3 mcmc samples with 2 warmup steps and 1 variational step - is this expected? forward_filter only takes about 0.6 seconds so this surprised me.

Here is the code - am I doing something stupid?

import tensorflow as tf
import tensorflow_probability as tfp
import time as tm

tf.random.set_seed(42)

tfd = tfp.distributions

num_observed_seasons = 10
num_seasons = 2
num_steps_per_season = 10
drift_scale = 1.0
observation_noise_scale = 0.2

initial_state_prior = tfd.MultivariateNormalLinearOperator(
    loc=tf.fill([num_seasons], 10.0),
    scale=tf.linalg.LinearOperatorScaledIdentity(num_seasons, 5.0),
)

drift_scale_prior = tfd.HalfNormal(scale=1.0)

num_timesteps = num_observed_seasons * num_steps_per_season

seasonal_ssm = tfp.sts.SeasonalStateSpaceModel(
    num_timesteps=num_timesteps,
    num_seasons=num_seasons,
    num_steps_per_season=num_steps_per_season,
    drift_scale=drift_scale,
    initial_state_prior=initial_state_prior,
    observation_noise_scale=observation_noise_scale,
    name="seasonal_ssm",
)

sampled_time_series = seasonal_ssm.sample()

start = tm.time()
lls, means, covs, *_ = seasonal_ssm.forward_filter(sampled_time_series)
end = tm.time()
print(f"finished forward_filter in {end - start:.4f} seconds!")

seasonal = tfp.sts.Seasonal(
    num_seasons=num_seasons,
    num_steps_per_season=num_steps_per_season,
    drift_scale_prior=drift_scale_prior,
    initial_effect_prior=initial_state_prior,
    name="seasonal",
)

start = tm.time()
mcmc, kernel_results = tfp.sts.fit_with_hmc(
    model=seasonal,
    observed_time_series=sampled_time_series,
    num_results=3,
    num_warmup_steps=2,
    num_variational_steps=1,
)
end = tm.time()
print(f"finished fit_with_hmc in {end - start:.4f} seconds!")

which outputs (alongside many warnings):

finished forward_filter in 0.6293 seconds!
finished fit_with_hmc in 133.7724 seconds!

Finally, in an attempt to speed this up, I tried wrapping tfp.sts.fit_with_hmc with @tf.function but this gives the following error:

@tf.function
def fit(model, observed_time_series):
    return tfp.sts.fit_with_hmc(seasonal_ssm, observed_time_series)


mcmc_tf_func, kernel_results_tf_func = fit(seasonal, sampled_time_series)

output:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/jeff/workspace/misc/fit_with_hmc_issue.py", line 66, in <module>
    mcmc_tf_func, kernel_results_tf_func = fit(seasonal, sampled_time_series)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 419, in __call__
    self._initialize(args, kwds, add_initializers_to=initializer_map)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 363, in _initialize
    *args, **kwds))
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1295, in _get_concrete_function_internal_garbage_collected
    graph_function, _, _ = self._maybe_define_function(args, kwargs)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1558, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1491, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py", line 692, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 316, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py", line 684, in wrapper
    ), args, kwargs)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 353, in converted_call
    result = converted_f(*effective_args, **kwargs)
  File "/tmp/tmpwthczynh.py", line 6, in tf__fit
    retval_ = ag__.converted_call('fit_with_hmc', tfp.sts, ag__.ConversionOptions(recursive=True, force_conversion=False, optional_features=(), internal_convert_user_code=True), (seasonal_ssm, observed_time_series), {})
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 260, in converted_call
    return _call_unconverted(f, args, kwargs)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 177, in _call_unconverted
    return f(*args, **kwargs)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow_probability/python/sts/fitting.py", line 495, in fit_with_hmc
    _, variational_distributions = make_variational()
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow/python/ops/template.py", line 380, in __call__
    return self._call_func(args, kwargs)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow/python/ops/template.py", line 343, in _call_func
    result = self._func(*args, **kwargs)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow_probability/python/sts/fitting.py", line 491, in make_variational
    init_batch_shape=chain_batch_shape, seed=seed())
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow_probability/python/sts/fitting.py", line 243, in build_factored_variational_loss
    q = _build_trainable_posterior(param, initial_loc_fn=initial_loc_fn)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow_probability/python/sts/fitting.py", line 68, in _build_trainable_posterior
    param.name + '_loc',
AttributeError: 'str' object has no attribute 'name'

originally defined at:
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow_probability/python/sts/fitting.py", line 494, in fit_with_hmc
    make_variational)
  File "/home/jeff/.virtualenvs/tf2/lib/python3.6/site-packages/tensorflow/python/ops/template.py", line 154, in make_template
    **kwargs)

Many thanks for any replies and all your amazing work on tensorflow probability!

Start in tensorflow_probability/python/sts/fitting.py at fit_with_hmc and compare its num_variational_steps default with the documented value. Reproduce the supplied Seasonal model timing and the @tf.function traceback, then trace the reported behavior. Done means the default discrepancy, slow runtime, and graph-mode error are each explained or resolved.