我尝试使用TensorFlow来测试卡尔曼滤波器，按照官方说明（https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/LinearGaussianStateSpaceModel）定义模型，生成样本，最后计算样本的对数似然值。
我正在运行指令提供的代码

import numpy as np
import matplotlib.pyplot as plt
import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
import matplotlib.pyplot as plt

tfd = tfp.distributions
ndims = 2
step_std = 1.0
noise_std = 5.0
model = tfd.LinearGaussianStateSpaceModel(
  num_timesteps=1000,
  transition_matrix=tf.linalg.LinearOperatorIdentity(ndims),
  transition_noise=tfd.MultivariateNormalDiag(
   scale_diag=step_std**2 * tf.ones([ndims])),
  observation_matrix=tf.linalg.LinearOperatorIdentity(ndims),
  observation_noise=tfd.MultivariateNormalDiag(
   scale_diag=noise_std**2 * tf.ones([ndims])),
  initial_state_prior=tfd.MultivariateNormalDiag(
   scale_diag=tf.ones([ndims])))

x = model.sample(1) # Sample from the prior on sequences of observations.
lp = model.log_prob(x) # Marginal likelihood of a (batch of) observations.
print(lp)

它需要30秒来计算对数似然。PS：我在colab上运行了代码，使用了GPU。
我的问题：为什么它这么慢，我怎样才能提高性能？
谢谢。