wrong batch_shape of distribution for observation distribution in HMM

[GongTen]

I got a serial of observations with shape of (N, 1). These observations correspond to two hidden states.

For hidden state 0, I declare the distribution as below:
dis0 = tfd.MixtureSameFamily( mixture_distribution=tfd.Categorical(probs=[0.4, 0.4, 0.2]), # mixture_distribution=tfd.Categorical(probs=[0.4, 0.4, 0.2]) components_distribution=tfd.Normal(loc=tf.Variable([100.0, 260.0, 180.0], trainable=True), scale=tf.Variable([10.0], trainable=True)) # tfd.Normal(loc=[100, 260, 180], scale=[10, 10, 10]) )
which mean it's a mixed Normal Distribution with three peaks.

For hidden. state 1, I declare the distribution as below:
dis1 = tfd.Normal(loc=tf.Variable(0.0, trainable=True), scale=tf.Variable(40.0, trainable=True))
which means it's a simple Normal Distribution

then I declare the 'observation distribution' as :
observation_distribution = tfd.Mixture( cat=tfd.Categorical(probs=[0.9, 0.1]), # 对应两个状态的概率 components=[dis_0, dis_1] )

After running the above code, I got observation_distribution.batch_shape = []

Even if I adjust the code as below:
`dis_swim = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(probs=[[0.4, 0.4, 0.2]]), # mixture_distribution=tfd.Categorical(probs=[0.4, 0.4, 0.2])
components_distribution=tfd.Normal(loc=tf.Variable([100.0, 260.0, 180.0], trainable=True), scale=tf.Variable([10.0], trainable=True)) # tfd.Normal(loc=[100, 260, 180], scale=[10, 10, 10])
)

dis_turn = tfd.Normal(loc=tf.Variable([0.0], trainable=True), scale=tf.Variable([40.0], trainable=True))
print(dis_turn.batch_shape, dis_turn.event_shape)

observation_distribution = tfd.Mixture(
cat=tfd.Categorical(probs=[[0.9, 0.1]]), # 对应两个状态的概率
components=[dis_swim, dis_turn]
)`

I got observation_distribution.batch_shape = [1]

!!!
However, the observation distribution in HMM needs batch_size equals 2 when I definate the HMM model as below:
hmm = tfd.HiddenMarkovModel( initial_distribution=tfd.Categorical(probs=[0.5, 0.5]), transition_distribution=tfd.Categorical(probs=[[0.9, 0.1], [0.1, 0.9]]), observation_distribution=observation_distribution, num_steps=N)

so I got the error "observation_distribution can't have scalar batches" or "transition_distribution and observation_distribution must agree on last dimension of batch size"

Is there anyone could introduct me to fix my code?

PS: Tensorflow:2.17.0 Tensorflow-probability 0.24.0