Add regression test for batched time_varying_P

Author: ricardoV94

tests/model/marginal/test_discrete_markov_chain.py

@@ -273,6 +273,43 @@ def test_marginalized_discrete_markov_chain_time_varying_P():
     np.testing.assert_allclose(marginal_m.compile_logp()({}), expected)
 
 
+def test_marginalized_discrete_markov_chain_batched_time_varying_P():
+    """Batched chains sharing one time-varying P. The batch axes of the chain must not collide
+    with P's core (time, from, to) axes in the forward filter."""
+    pi0 = np.array([0.6, 0.4])
+    sigma = 0.8
+    obs = np.array([[0.1, 1.2, 0.8, -0.2], [0.9, -0.1, 0.2, 1.1], [0.5, 0.7, 1.0, 0.3]])
+    n_chains, T, k = *obs.shape, 2
+    A_t = np.array(
+        [
+            [[0.9, 0.1], [0.2, 0.8]],
+            [[0.4, 0.6], [0.3, 0.7]],
+            [[0.1, 0.9], [0.5, 0.5]],
+        ]
+    )  # (T - 1, k, k), no batch axis: every chain shares the same per-step kernels
+
+    with pm.Model() as m:
+        init = pm.Categorical.dist(p=pi0, shape=(n_chains,))
+        states = DiscreteMarkovChain(
+            "states", P=A_t, init_dist=init, time_varying_P=True, shape=(n_chains, T)
+        )
+        pm.Normal("emission", mu=states, sigma=sigma, observed=obs)
+
+    marginal_m = marginalize(m, [states])
+
+    # The chains are independent, so brute force each one over its k**T paths and sum.
+    expected = 0.0
+    for b in range(n_chains):
+        log_joint = []
+        for s in itertools.product(range(k), repeat=T):
+            lp = np.log(pi0[s[0]]) + sum(np.log(A_t[t - 1, s[t - 1], s[t]]) for t in range(1, T))
+            lp += norm.logpdf(obs[b], loc=np.array(s), scale=sigma).sum()
+            log_joint.append(lp)
+        expected += logsumexp(log_joint)
+
+    np.testing.assert_allclose(marginal_m.compile_logp()({}), expected)
+
+
 # The higher-order (n_lags > 1) forward scan currently OOMs under the default numba backend
 # (a numba fast_run scan-lowering bug, not a logic issue; see numba_nlags2_repro.py). Compile the
 # logp under the C VM so these still verify correctness; drop the override once numba is fixed.