Poor fit with Linear Regression in tfp.sts

Hi,
I am trying to fit a time series model with exogenous variable. The linear regression part doesn't seem to be able to find weights correctly. Here is a simple example to reproduce the problem-

import numpy as np
import tensorflow_probability as tfp
import tensorflow as tf
from sklearn.preprocessing import StandardScaler
weights=np.random.randn(5)+np.arange(1,6)
design_matrix=np.random.normal(4,2,(200,5))
tseries = np.matmul(design_matrix,weights)
model=tfp.sts.Sum([tfp.sts.LinearRegression(design_matrix=design_matrix)],observed_time_series=tseries)
variational_posterior = tfp.sts.build_factored_surrogate_posterior(model)
optimizer = tf.optimizers.Adam(learning_rate=0.1)
tfp.vi.fit_surrogate_posterior(target_log_prob_fn=model.joint_log_prob(observed_time_series=tseries),
                                                     surrogate_posterior=variational_posterior,optimizer=optimizer,num_steps=200
                                                     )

samples = variational_posterior.sample(50)
fitted_weights=np.mean(samples['LinearRegression/_weights'], axis=0)
error_rms = np.sqrt(np.mean(np.power(weights-fitted_weights,2)))
error_rms
3.02

if I normalize the design_matrix, the numbers look much better

normalizer = StandardScaler()
normalized_matrix = normalizer.fit_transform(design_matrix)
model=tfp.sts.Sum([tfp.sts.LinearRegression(design_matrix=normalized_matrix)],observed_time_series=tseries)
variational_posterior = tfp.sts.build_factored_surrogate_posterior(model)
optimizer = tf.optimizers.Adam(learning_rate=0.1)
tfp.vi.fit_surrogate_posterior(target_log_prob_fn=model.joint_log_prob(observed_time_series=tseries),
                                                     surrogate_posterior=variational_posterior,optimizer=optimizer,num_steps=200
                                                     )

samples = variational_posterior.sample(50)
fitted_weights=np.mean(samples['LinearRegression/_weights'], axis=0)/normalizer.scale_
error_rms = np.sqrt(np.mean(np.power(weights-fitted_weights,2)))
error_rms
0.56

Start by running the Python reproduction using tfp.sts.LinearRegression, tfp.sts.Sum, build_factored_surrogate_posterior, and tfp.vi.fit_surrogate_posterior. Compare the raw and StandardScaler-normalized design-matrix cases, then inspect the LinearRegression weight estimates; done means the reported fitting discrepancy is understood and the raw-input behavior is corrected or documented with a regression check.