Python Numpyro AR(1)均值切换模型抽样不一致性

我试图根据潜在状态S = 0，1来估计具有切换平均值的AR(1)过程y，该过程y以固定转移概率的马尔可夫过程(如here)的形式发展.简而言之，它采取以下形式:

y_t - mu_{0/1} = phi * (y_{t-1} - mu_{0/1})+ epsilon_t

其中，如果STATE_T=0，则使用mu_0，如果STATE_T=1，则使用MU_1. 我正在使用带有DiscreteHMCGibbs的JAX/Numpyro(尽管带有潜伏态枚举的普通坚果会产生相同的结果)，但我似乎无法使采样器正常工作.从我运行的所有诊断来看，似乎所有的超参数都停留在初始化值上，摘要相应地返回所有std==0.下面我有一个重现我的问题的MWE.我在实现过程中犯了什么明显的错误吗？

MWe:

import jax.numpy as jnp
import numpyro
import numpyro.distributions as dist
from numpyro.contrib.control_flow import scan
from numpyro.infer import MCMC, NUTS,DiscreteHMCGibbs
from jax import random, pure_callback
import jax
import numpy as np

def generate_synthetic_data(T=100, mu=[0, 5], phi=0.5, sigma=1.0, p=np.array([[0.95, 0.05], [0.1, 0.9]])):
    states = np.zeros(T, dtype=np.int32)
    y = np.zeros(T)
    current_state = np.random.choice([0, 1], p=[0.5, 0.5])
    states[0] = current_state
    y[0] = np.random.normal(mu[current_state], sigma)

    for t in range(1, T):
        current_state = np.random.choice([0, 1], p=p[current_state,:])
        states[t] = current_state
        y[t] = np.random.normal(mu[current_state] + phi * (y[t-1] - mu[current_state]), sigma)

    return y, states


def mean_switching_AR1_model(y):
    T = len(y)
    phi = numpyro.sample('phi', dist.Normal(0, 1))
    sigma = numpyro.sample('sigma', dist.Exponential(1))
    
    
    with numpyro.plate('state_plate', 2):
        mu = numpyro.sample('mu', dist.Normal(0, 5))
        p = numpyro.sample('p', dist.Dirichlet(jnp.ones(2)))

    probs_init = numpyro.sample('probs_init', dist.Dirichlet(jnp.ones(2)))
    s_0 = numpyro.sample('s_0', dist.Categorical(probs_init))

    def transition_fn(carry, y_t):
        prev_state = carry
        state_probs = p[prev_state]
        state = numpyro.sample('state', dist.Categorical(state_probs))

        mu_state = mu[state]
        y_mean = mu_state + phi * (y_t - mu_state)
        y_next = numpyro.sample('y_next', dist.Normal(y_mean, sigma), obs=y_t)
        return state, (state, y_next)

    _ , (signal, y)=scan(transition_fn, s_0, y[:-1], length=T-1)
    return (signal, y)

# Synthetic data generation
T = 1000
mu_true = [0, 3]
phi_true = 0.5
sigma_true = 0.25
transition_matrix_true = np.array([[0.95, 0.05], [0.1, 0.9]])
y, states_true = generate_synthetic_data(T, mu=mu_true, phi=phi_true, sigma=sigma_true, p=transition_matrix_true)


rng_key = random.PRNGKey(0)
nuts_kernel = NUTS(mean_switching_AR1_model)
gibbs_kernel = DiscreteHMCGibbs(nuts_kernel, modified=True)

# Run MCMC
mcmc = MCMC(gibbs_kernel, num_samples=1000, num_warmup=1000)
mcmc.run(rng_key, y=y)
mcmc.print_summary()