bayespy.nodes.Choose

bayespy.nodes.Choose(z, *nodes)

Choose plate elements from nodes based on a categorical variable.

For instance:

>>> import bayespy as bp
>>> z = [0, 0, 2, 1]
>>> x0 = bp.nodes.GaussianARD(0, 1)
>>> x1 = bp.nodes.GaussianARD(10, 1)
>>> x2 = bp.nodes.GaussianARD(20, 1)
>>> x = bp.nodes.Choose(z, x0, x1, x2)
>>> print(x.get_moments()[0])
[ 0.  0. 20. 10.]

This is basically just a thin wrapper over applying Gate node over the concatenation of the nodes.