import aesara.tensor as at

from aesara.ifelse import ifelse

from aeppl.joint_logprob import factorized_joint_logprob


srng = at.random.RandomStream(seed=2320)

I_rv = srng.bernoulli(0.5, name="I")
X_rv = srng.normal(0, 1, name="X")
Y_rv = srng.gamma(0.5, 0.5, name="Y")

Z_rv = ifelse(I_rv, X_rv, Y_rv)
Z_rv.name = "Z"

z_vv = Z_rv.clone()
i_vv = I_rv.clone()
logp_parts = factorized_joint_logprob({Z_rv: z_vv, I_rv: i_vv})