Hello, I read several posts and tried different approaches before writing this post, but I have been stuck for a moment now and I cannot make a simple time series work when the variables have more than one dimension.

Below is an example of a simple HMM with a latent variable Z that evolves over time with Gaussian noise, and another variable X (observable) with respective emission noise.

```
T = 50 # number of time steps
F = 2 # number of features (dimension of each state/observation)
sd_z = 1 # transition noise (single noise for all features)
sd_x = 1 # emission noise (single noise for all features)
# Data generation
np.random.seed(0)
x = np.zeros((T, F))
z = np.zeros((T, F))
for t in range(T):
if t == 0:
z[0] = norm(loc=0, scale=1).rvs(size=F)
else:
z[t] = norm(loc=z[t - 1], scale=sd_z).rvs()
x[t] = norm(loc=z[t], scale=sd_x).rvs()
# Model definition and posterior inference
coords = {"time": np.arange(T), "features": np.arange(F)}
hmm = pm.Model(coords=coords)
with hmm:
x_obs = pm.MutableData("x_obs", x, dims=("time", "features"))
# Standard deviation of the transition and emission distributions
i_sd_z = pm.HalfNormal(name="i_sd_z", sigma=1)
i_sd_x = pm.HalfNormal(name="i_sd_x", sigma=1)
# Transition distribution of latent variable Z
i_z = pm.GaussianRandomWalk(name="i_z", init_dist=pm.Normal.dist(mu=0, sigma=1, shape=(1, F)),
mu=0, sigma=i_sd_z, dims=("time", "features"))
# Emission distribution of the observed variable X
i_x = pm.Normal(name="i_x", mu=i_z, sigma=i_sd_x, observed=x_obs)
idata = pm.sample(1000, init="adapt_diag", tune=1000, chains=2, random_seed=0)
# Plotting results
az.plot_trace(idata, var_names=["i_sd_z", "i_sd_x"])
plt.show()
m1 = idata.posterior["i_z"].sel(features=0).mean(dim=["chain", "draw"])
sd1 = idata.posterior["i_z"].sel(features=0).std(dim=["chain", "draw"])
plt.figure(figsize=(15, 8))
plt.plot(range(T), z[:, 0], label="Real", color="tab:blue", marker="o")
plt.plot(range(T), m1, label="Inferred", color="tab:orange", marker="o")
plt.fill_between(range(T), m1 - sd1, m1 + sd1, color="tab:orange", alpha=0.4)
plt.legend()
plt.show()
```

See from the plots below that the inferred values of the standard deviations are not correct as well as the samples of the latent variable Z. The problem persists even if I set F to 1. Although, if I remove the â€śfeaturesâ€ť dimension, everything works well so it seems the problem is with the indexing inside the GaussianRandomWalk distribution.

Python version: 3.9

PyMC version: 5.0.1