STS models running slower on GPUs

This simple example running on Colab shows this issue:

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


tfd = tfp.distributions
tfb = tfp.bijectors

ds = np.random.rand(1300, 5)

components = []
ots = ds[:1000, 0].astype(np.float32)

obs_sd = ots.std()

level_component = tfp.sts.LocalLevel(observed_time_series=ots)
components.append(level_component)

linear_component = tfp.sts.LinearRegression(design_matrix=ds[:,1:].astype(np.float32))
components.append(linear_component)

model = tfp.sts.Sum(components, observed_time_series=ots)

t0 = time()
optimizer = tf.optimizers.Adam(learning_rate=0.1)
variational_steps = 200
variational_posteriors = tfp.sts.build_factored_surrogate_posterior(model=model)

@tf.function()
def _run_vi():
    tfp.vi.fit_surrogate_posterior( 
        target_log_prob_fn=model.joint_log_prob(
            observed_time_series=ots
        ),
        surrogate_posterior=variational_posteriors,
        optimizer=optimizer,
        num_steps=variational_steps
    )

    samples = variational_posteriors.sample(100)
    return samples, None
_run_vi()
t1 = time()

print(t1 - t0)

Running on CPU takes a few seconds whereas running on GPU takes minutes. I also tested the same procedure using TFP's example notebook on Colab and again running on GPU also took longer (between 2~3x).

I tried testing with bigger datasets to see if data volume was the issue but as I increased it ten fold the GPU could no longer finish its process on Colab.

Also tested on previous versions of Tensorflow and Probability but the issue remained. Is there something that changed that made GPUs slower?

Thanks in advance!

Start by reproducing the provided Colab example and timing the _run_vi entry point on CPU and GPU. Compare the result with the linked Structural_Time_Series_Modeling_Case_Studies_Atmospheric_CO2_and_Electricity_Demand.ipynb, including the reported TensorFlow and TensorFlow Probability versions. Done means identifying a reproducible reason for the GPU slowdown or documenting the remaining evidence and affected configuration.