Unable to get DiscreteMarkovChain to work in PyMC

I am trying to fit a two-state regime-switching model using PyMC.

The full model is given below. The exact details of the model are not important (if you're curious anyways, I'm trying to fit a geometric Brownian motion whose coefficients μ and σ are functions of the regime).

I get a nebulous error I am unable to interpret (also provided below). My question is as follows: am I using DiscreteMarkovChain or is there a bug? Is there a better/more idiomatic way to express this model?

with pm.Model() as model:
    π0 = pm.Uniform("π01", lower=0., upper=10.)
    π1 = pm.Uniform("π10", lower=0., upper=10.)
    P = pm.math.stack([[1. - π0 * Δt, π0 * Δt], [π1 * Δt, 1. - π1 * Δt]])
    state = pymc_experimental.distributions.timeseries.DiscreteMarkovChain("state", P=P, steps=Δx.size - 1)

    μ0 = pm.Normal("μ0", mu=0, sigma=0.1)
    μ1 = pm.Normal("μ1", mu=0, sigma=0.1)
    
    σ0 = pm.HalfNormal("σ0", sigma=0.1)
    σ1 = pm.HalfNormal("σ1", sigma=0.1)
    
    μ = pm.math.switch(state, μ0, μ1)
    σ = pm.math.switch(state, σ0, σ1)
 
    ΔX = pm.Normal("ΔX", mu=(μ - σ**2 / 2) * Δt, sigma=σ * math.sqrt(Δt), observed=Δx)

Traceback (most recent call last):
  File "/home/parsiad/venv/lib/python3.12/site-packages/pymc/sampling/parallel.py", line 128, in run
    self._start_loop()
  File "/home/parsiad/venv/lib/python3.12/site-packages/pymc/sampling/parallel.py", line 180, in _start_loop
    point, stats = self._step_method.step(self._point)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/parsiad/venv/lib/python3.12/site-packages/pymc/step_methods/compound.py", line 232, in step
    point, sts = method.step(point)
                 ^^^^^^^^^^^^^^^^^^
  File "/home/parsiad/venv/lib/python3.12/site-packages/pymc/step_methods/arraystep.py", line 101, in step
    apoint, stats = self.astep(q)
                    ^^^^^^^^^^^^^
  File "/home/parsiad/venv/lib/python3.12/site-packages/pymc/step_methods/metropolis.py", line 274, in astep
    accept_rate = self.delta_logp(q, q0d)
                  ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/parsiad/venv/lib/python3.12/site-packages/pytensor/compile/function/types.py", line 972, in __call__
    raise_with_op(
  File "/home/parsiad/venv/lib/python3.12/site-packages/pytensor/link/utils.py", line 528, in raise_with_op
    raise exc_value.with_traceback(exc_trace)
  File "/home/parsiad/venv/lib/python3.12/site-packages/pytensor/compile/function/types.py", line 959, in __call__
    self.vm()
  File "/home/parsiad/venv/lib/python3.12/site-packages/pytensor/graph/op.py", line 524, in rval
    r = p(n, [x[0] for x in i], o)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/parsiad/venv/lib/python3.12/site-packages/pytensor/tensor/subtensor.py", line 2754, in perform
    rval = inputs[0].__getitem__(tuple(inputs[1:]))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
IndexError: index 2 is out of bounds for axis 0 with size 2
Apply node that caused the error: AdvancedSubtensor(Join.0, Subtensor{start:stop}.0, Subtensor{start:}.0)
Toposort index: 32
Inputs types: [TensorType(float64, shape=(2, 2)), TensorType(int64, shape=(24113,)), TensorType(int64, shape=(24113,))]
Inputs shapes: [(2, 2), (24113,), (24113,)]
Inputs strides: [(16, 8), (8,), (8,)]
Inputs values: [array([[0.98880855, 0.01119145],
       [0.02356341, 0.97643659]]), 'not shown', 'not shown']
Inputs type_num: [12, 7, 7]
Outputs clients: [[CAReduce{Composite{(i0 + log(i1))}, axes=None}(AdvancedSubtensor.0)]]

I searched the PyMC documentation for alternatives to DiscreteMarkovChain. I also tried to debug my current code.

I believe that you should go with smth like this. The main idea is to sample state variable using BinaryGibbsMetropolis. It is also better to avoid running 1. - π1 * Δt because it might create negative probabilities.

import pymc as pm
from pymc_experimental.distributions import DiscreteMarkovChain
import math 
import numpy as np



with pm.Model() as model:

    Δx = np.random.randn(20)
    Δt =1

    π0 = pm.Uniform("π01", lower=0., upper=1.)
    π1 = pm.Uniform("π10", lower=0., upper=1.)
    s0 = pm.Bernoulli.dist(p=0.5)
    P = pm.Dirichlet("P", a=[π0 * Δt, π1 * Δt], size=(2,))
    state = DiscreteMarkovChain("state", P=P,init_dist=s0, steps=Δx.size-1)

    μ0 = pm.Normal("μ0", mu=0, sigma=0.1)
    μ1 = pm.Normal("μ1", mu=0, sigma=0.1)
    
    σ0 = pm.HalfNormal("σ0", sigma=0.1)
    σ1 = pm.HalfNormal("σ1", sigma=0.1)
    
    μ = pm.math.switch(state, μ0, μ1)
    σ = pm.math.switch(state, σ0, σ1)
 
    ΔX = pm.Normal("ΔX", mu=(μ - σ**2 / 2) * Δt, sigma=σ * math.sqrt(Δt), observed=Δx)


    idata = pm.sample(
        step=[
            pm.BinaryGibbsMetropolis([state]),
        ],
        draws=5_000,
    )