def to_child(self):
mean = self.parent / self.message_to_parent
m = mean.mean()
v = mean.log_var() + self.variance.value()
message_to_child = GaussianArray.from_array(m, v)
self.child.update(self.message_to_child, message_to_child)
self.message_to_child = message_to_child
def to_sum(self, x):
x = x / self.message_to_x
self.message_to_sum = GaussianArray.from_array(
tf.math.reduce_sum(x.mean(), -1),
tf.math.reduce_sum(x.log_var(), -1)
)
return self.message_to_sum
def to_parent(self):
x = self.child / self.message_to_child
m = x.mean()
v = x.variance_safe() + self.variance.value()
message_to_parent = GaussianArray.from_array(m, v)
self.parent.update(self.message_to_parent, message_to_parent)
self.message_to_parent = message_to_parent
def to_child(self):
x = self.parent / self.message_to_parent
message_to_child = GaussianArray.from_array(
tf.math.reduce_sum(x.mean(), -1),
tf.math.reduce_sum(x.log_var(), -1)
)
self.child.update(self.message_to_child, message_to_child)
self.message_to_child = message_to_child
def to_child(self):
x = self.parent / self.message_to_parent
weight = self.weight.value()
bias = self.bias.value()
m = tf.tensordot(x.mean(), weight, 1) + bias
v = tf.tensordot(x.log_var(), weight ** 2, 1)
message_to_child = GaussianArray.from_array(m, v)
self.child.update(self.message_to_child, message_to_child)
self.message_to_child = message_to_child
def to_x(self, x, A):
# compute incoming message
from_x = x / self.message_to_x
# add noise
from_x = GaussianArray.from_array(
mean=from_x.mean(),
variance=from_x.log_var() + self.variance.value()
)
# compute outgoing message
A = A.proba() # TODO this assumes 0/1, need to fix later
stnr = from_x.mean() * tf.math.sqrt(from_x.precision()) * tf.cast(2*A - 1, tf.float32)
vf = tfp.distributions.Normal(0., 1.).prob(stnr) / tfp.distributions.Normal(0., 1.).cdf(stnr)
wf = vf * (stnr + vf)
m = from_x.mean() + tf.math.sqrt(from_x.log_var()) * vf * tf.cast(2 * A - 1, tf.float32)
v = from_x.log_var() * (1. - wf)
# add noise
v += self.variance.value()
self.message_to_x = GaussianArray.from_array(m, v)
return self.message_to_x
def to_product(self, x):
x = x / self.message_to_x
m0 = tf.expand_dims(x.mean(), 0)
m1 = tf.expand_dims(x.mean(), 1)
v0 = tf.expand_dims(x.log_var(), 0)
v1 = tf.expand_dims(x.log_var(), 1)
m = m0 * m1
v = m0 ** 2 * v1 + m1 ** 2 * v0 + v0 * v1
self.message_to_product = GaussianArray.from_array(m, v)
return self.message_to_product
def to_result(self, x):
# compute incoming message
from_x = x / self.message_to_x
# add noise
from_x = GaussianArray.from_array(
mean=from_x.mean(),
variance=from_x.log_var() + self.variance.value()
)
# compute probability
proba = 1. - tfp.distributions.Normal(*from_x.mean_and_variance()).cdf(0.0)
return BernoulliArray.from_array(proba)
def to_x(self, x, A):
# TODO: assumes 0/1 only, need to fix later with missing values
x = x / self.message_to_x
A = A.proba()
stnr = x.mean() * tf.math.sqrt(x.precision()) * tf.cast(2*A - 1, tf.float32)
vf = tfp.distributions.Normal(0., 1.).prob(stnr) / tfp.distributions.Normal(0., 1.).cdf(stnr)
wf = vf * (stnr + vf)
m = x.mean() + tf.math.sqrt(x.log_var()) * vf * tf.cast(2 * A - 1, tf.float32)
v = x.log_var() * (1. - wf)
self.message_to_x = GaussianArray.from_array(m, v)
return self.message_to_x
def to_x(self, x, sum):
sum = sum / self.message_to_sum
x = x / self.message_to_x
m = tf.expand_dims(sum.mean(), -1) - tf.math.reduce_sum(x.mean(), -1, keepdims=True) + x.mean()
v = tf.where(
x.is_uniform(),
np.inf,
tf.expand_dims(sum.log_var(), -1) + tf.math.reduce_sum(x.log_var(), -1, keepdims=True) - x.log_var()
)
self.message_to_x = GaussianArray.from_array(m, v)
return self.message_to_x
def __init__(self, child: GaussianArray, mean: float = 0., variance: float = 1., initial=None, name=""):
super().__init__()
self._deterministic = False
self.mean = ParameterArray(mean, True, name=name+".mean")
self.variance = ParameterArrayLogScale(variance, True, name=name+".variance")
self._parameters = {"mean": self.mean, "variance": self.variance}
self.child = child
self.shape = child.shape()
self.prior = GaussianArray.from_shape(self.shape, self.mean.value(), self.variance.value())
# initialize child
self.message_to_child = GaussianArray.from_array(initial, tf.ones_like(initial) * variance * 0.1)
self.child.set_to(self.message_to_child)
def to_parent(self):
s = self.child / self.message_to_child
x = self.parent / self.message_to_parent
m = tf.expand_dims(s.mean(), -1) - tf.math.reduce_sum(x.mean(), -1, keepdims=True) + x.mean()
v = tf.where(
x.is_uniform(),
np.inf,
tf.expand_dims(s.log_var(), -1) + tf.math.reduce_sum(x.log_var(), -1, keepdims=True) - x.log_var()
)
message_to_parent = GaussianArray.from_array(m, v)
self.parent.update(self.message_to_parent, message_to_parent)
self.message_to_parent = message_to_parent
def to_elbo(self, x, A):
mean = x
x = GaussianArray.from_array(
mean=x.mean(),
variance=x.log_var() + self.variance.value()
)
elbo = mean.log_var() + mean.mean() ** 2
elbo += x.log_var() + x.mean() ** 2
elbo += -2. * x.mean() * mean.mean()
elbo /= self.variance.value()
# elbo += tf.math.log(2 * np.pi)
elbo += tf.math.log(self.variance.value())
return -.5 * tf.reduce_sum(elbo)
def to_parent(self):
x = self.parent / self.message_to_parent
A = self.child.proba()
A_safe = tf.where(tf.math.is_nan(A), 0.5, A)
stnr = x.mean() * tf.math.sqrt(x.precision()) * tf.cast(2 * A_safe - 1, tf.float32)
vf = tfp.distributions.Normal(0., 1.).prob(stnr) / tfp.distributions.Normal(0., 1.).cdf(stnr)
wf = vf * (stnr + vf)
m = x.mean() + tf.math.sqrt(x.log_var()) * vf * tf.cast(2 * A_safe - 1, tf.float32)
m = tf.where(A_safe == 0.5, 0., m)
v = x.log_var() * (1. - wf)
v = tf.where(A_safe == 0.5, 1.0e10, v)
message_to_parent = GaussianArray.from_array(m, v)
self.parent.update(self.message_to_parent, message_to_parent)
self.message_to_parent = message_to_parent
def to_child(self):
p, mtp = self.parent.natural()
p0 = tf.expand_dims(p, 0)
p1 = tf.expand_dims(p, 1)
mtp0 = tf.expand_dims(mtp, 0)
mtp1 = tf.expand_dims(mtp, 1)
v0 = 1. / p0
v1 = 1. / p1
m0 = mtp0 * v0
m1 = mtp1 * v1
m = m0 * m1
v = m0 ** 2 * v1 + m1 ** 2 * v0 + v0 * v1
# new marginal
child = GaussianArray.from_array(m, v)
# compute what should be the message
message_to_child = (child * self.message_to_child) / self.child
# update and store
self.child.update(self.message_to_child, message_to_child)
self.message_to_child = message_to_child
def to_x(self, mean, variance):
mean = mean / self.message_to_mean
m = mean.mean()
v = mean.log_var() + variance
self.message_to_x = GaussianArray.from_array(m, v)
return self.message_to_x
def to_result(self, x, weight, bias):
x = x / self.message_to_x
m = tf.tensordot(x.mean(), weight, 1) + bias
v = tf.tensordot(x.log_var(), weight ** 2, 1)
self.message_to_result = GaussianArray.from_array(m, v)
return self.message_to_result
def to_mean(self, x, variance):
x = x / self.message_to_x
m = x.mean()
v = x.log_var() + variance
self.message_to_mean = GaussianArray.from_array(m, v)
return self.message_to_mean
def to_mean(self, x, variance):
m = x.mean()
v = variance * tf.ones_like(m) + x.variance_safe() # maybe this contains infinity
self.message_to_mean = GaussianArray.from_array(m, v)
return self.message_to_mean
def _break_symmetry(self):
self.factors["latent_prior"].message_to_x = GaussianArray.from_array(
mean=tf.random.normal((self.N, self.K), 0., 1.),
variance=tf.ones((self.N, self.K)) * 1000.
)
