Discrete Bayesian network on Tensorflow Probability, Edward2, and Python

Question

I have a simple Bayesian network:

        state
      /       \
     /         \
    V           V  
signal_1    signal_2

with random variables "state", "signal_1", and "signal_2" with corresponding discrete values: Val(state) = {0, 1, 2, 3}, Val(signal_1) = {0, 1}, and Val(signal_2) = {0, 1}.

I have marginal probability distribution P(state) and conditional probablity distributions P(signal_1|state) and P(signal_2|state) as tables.

Joint probability P(state, signal_1, signal_2) is equal to P(state) * P(signal_1|state) * P(signal_2|state) and log P(state, signal_1, signal_2) = log P(state) + log P(signal_1|state) + log P(signal_2|state).

I'm trying to construct this model in TensorFlow Probability:

Variant 1:

import tensorflow as tf
tfd = tfp.distributions

def joint_log_prob(state, signal_1, signal_2):
        state_prob = tf.constant([0.5, 0.1, 0.3, 0.1])
        signal_1_prob = tf.constant([[0.03, 0.97], [0.05, 0.95], [0.20, 0.80], [0.97, 0.03]])
        signal_2_prob = tf.constant([[0.03, 0.97], [0.97, 0.03], [0.97, 0.03], [0.99, 0.01]])

        state_dist = tfd.Categorical(probs=state_prob)
        signal_1_dist = tfd.Categorical(probs=signal_1_prob[state]) # state=0,1,2,3
        signal_2_dist = tfd.Categorical(probs=signal_2_prob[state]) # state=0,1,2,3

        return state_dist.log_prob(state) + signal_1_dist.log_prob(signal_1) + signal_2_dist.log_prob(signal_2)

My evidence is, for example, "signal_1 = 1" and "signal_2 = 0" and I want to get a probability of "state = 0", i.e. P(state=0|signal_1=1, signal_2=0)

I define the function:

def unnormalized_log_posterior(state):
    return joint_log_prob(state, 1, 0)

Variant 2 via Edward2:

import tensorflow as tf
from tensorflow_probability import edward2 as ed

def model():
    state_prob = tf.constant([0.5, 0.1, 0.3, 0.1])
    signal_1_prob = tf.constant([[0.03, 0.97], [0.05, 0.95], [0.20, 0.80], [0.97, 0.03]])
    signal_2_prob = tf.constant([[0.03, 0.97], [0.97, 0.03], [0.97, 0.03], [0.99, 0.01]])

    state = ed.Categorical(probs=state_prob, name='state')
    signal_1 = ed.Categorical(probs=signal_1_prob[state], name='signal_1')
    signal_2 = ed.Categorical(probs=signal_2_prob[state], name='signal_2')

    return [state, signal_1, signal_2] # what does the function have to return? [state, signal_1, signal_2]?

Then I get the joint log-probability function

log_joint = ed.make_log_joint_fn(model) # joint log-probability function

which can be used via lambda-expression:

lambda x: log_joint(state=x, signal_1=1, signal_2=0)

Did I define the model correctly?

The question is how do I continue to use MCMC? Which MCMC method can I use for this discrete model? Can someone show a MCMC code for my model? How can I tell the MCMC algorithm to take samples for "state" only from the set {0, 1, 2, 3}?

Thank you!