[Python] How can I run Bayesian Network by pomegrante?

I want to calculate the probability by using Bayesian Network on pomegrante. Here is the example of the code in tutorial:

from pomegranate.distributions import Categorical
from pomegranate.distributions import ConditionalCategorical
from pomegranate.bayesian_network import BayesianNetwork

d1 = Categorical([[0.1, 0.9]])
d2 = ConditionalCategorical([[[0.4, 0.6], [0.3, 0.7]]])

model = BayesianNetwork([d1, d2], [(d1, d2)])
X = numpy.random.randint(2, size=(10, 2))
model.probability(X)


I think I do the same things in tutorial:

import torch
from pomegranate import bayesian_network, distributions
import numpy as np

# Create distribution node
rain = distributions.Categorical(
probs= torch.tensor([[
0.7,
0.2,
0.1
]])
)

maintances = distributions.ConditionalCategorical(
probs= torch.tensor([[
[0.4, 0.6],
[0.2, 0.8],
[0.1, 0.9]
]])
)

train = distributions.ConditionalCategorical(
probs= torch.tensor([[
[0.8, 0.2],
[0.9, 0.1],
[0.6, 0.4],
[0.7, 0.3],
[0.4, 0.6],
[0.5, 0.5]
]])
)

appointment = distributions.ConditionalCategorical(
probs= torch.tensor([[
[0.9, 0.1],
[0.6, 0.4]
]])
)



# Create a Bayersian Network and add states
model = bayesian_network.BayesianNetwork()
model.add_distributions([rain, maintances, train, appointment])

# Add edge
model.add_edge(rain, maintances)
model.add_edge(rain, train)
model.add_edge(maintances, train)
model.add_edge(train, appointment)

X = np.column_stack([
np.random.randint(0, 3, size=10), # rain: 3 categories
np.random.randint(0, 2, size=10), # maintances: 2 categories
np.random.randint(0, 2, size=10), # train: 2 categories
np.random.randint(0, 2, size=10) # appointment: 2 categories
])

print(model.probability(X))


Here is the problem when I run the code: Traceback (most recent call last):

File "/home/crush_dpl/cs50_ai/week2/src2/learn.py", line 60, in print(model.probability(X)) ^^^^^^^^^^^^^^^^^^^^ File "/home/crush_dpl/.local/lib/python3.12/site-packages/pomegranate/distributions/distribution.py", line 69, in probability return torch.exp(self.log_probability(X)) ^^^^^^^^^^^^^^^^^^^^^^^ File "/home/crush_dpl/.local/lib/python3.12/site-packages/pomegranate/bayesian_network.py", line 354, in log_probability logps += distribution.log_probability(X) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/crush_dpl/.local/lib/python3.12/site-packages/pomegranate/distributions/conditional_categorical.py", line 210, in log_probability X = _check_parameter(_cast_as_tensor(X), "X", ndim=3, ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/home/crush_dpl/.local/lib/python3.12/site-packages/pomegranate/_utils.py", line 250, in _check_parameter raise ValueError("Parameter {} must have shape {}".format( ValueError: Parameter X must have shape (-1, 2, 1)

Continue reading...