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. )