def __call__(self, x): x_ad = self.adapter['in']['scale'] * x + self.adapter['in']['shift'] self.pre_adapt = self.mlp(x_ad) mean = self.adapter['out'][ 'scale'] * self.pre_adapt.mean + self.adapter['out']['shift'] cov = tf.reshape(self.adapter['out']['scale'], [-1, 1]) * tf.reshape( self.adapter['out']['scale'], [1, -1]) * self.pre_adapt.var return gv.GaussianVar(mean, cov)
def __call__(self, x): A = [aa.value.mean for aa in self.A] b = [bb.value.mean for bb in self.b] self.h = [x] self.h.append(tf.matmul(x, A[0]) + b[0]) for L in range(1, len(A)): self.h.append(tf.matmul(tf.nn.relu(self.h[-1]), A[L]) + b[L]) return gv.GaussianVar(self.h[-1], tf.constant(0.0, dtype=tf.float32))
def linear(x, A, b): """ compute y = x^T A + b """ x_mean = x.mean y_mean = tf.matmul(x_mean, A.mean) + b.mean x_cov = x.var y_cov = linear_covariance(x_mean, x_cov, A, b) return gv.GaussianVar(y_mean, y_cov)
def linear_certain_activations(x_certain, A, b): """ compute y = x^T A + b assuming x has zero variance """ x_mean = x_certain xx = x_mean * x_mean y_mean = tf.matmul(x_mean, A.mean) + b.mean y_cov = tf.matrix_diag(tf.matmul(xx, A.var) + b.var) return gv.GaussianVar(y_mean, y_cov)
def linear_relu_diagonal(x, A, b): """ compute y = relu(x)^T A + b """ x_var_diag = x.var sqrt_x_var_diag = tf.sqrt(x_var_diag) mu = x.mean / (sqrt_x_var_diag + EPSILON) pdf = bu.standard_gaussian(mu) cdf = bu.gaussian_cdf(mu) softrelu = pdf + mu * cdf z_mean = sqrt_x_var_diag * softrelu y_mean = tf.matmul(z_mean, A.mean) + b.mean z_var = x_var_diag * (cdf + mu * softrelu - tf.square(softrelu)) y_cov = linear_covariance_diagonal(z_mean, z_var, A, b) return gv.GaussianVar(y_mean, y_cov)
def linear_heaviside(x, A, b): """ compute y = heaviside(x)^T A + b """ x_var_diag = tf.matrix_diag_part(x.var) mu = x.mean / (tf.sqrt(x_var_diag) + EPSILON) def heaviside_covariance(x): mu1 = tf.expand_dims(mu, 2) mu2 = tf.transpose(mu1, [0, 2, 1]) s11s22 = tf.expand_dims(x_var_diag, axis=2) * tf.expand_dims( x_var_diag, axis=1) rho = x.var / (tf.sqrt(s11s22)) # + EPSILON) rho = tf.clip_by_value(rho, -1 / (1 + EPSILON), 1 / (1 + EPSILON)) return bu.heavy_g(rho, mu1, mu2) z_mean = bu.gaussian_cdf(mu) y_mean = tf.matmul(z_mean, A.mean) + b.mean z_cov = heaviside_covariance(x) y_cov = linear_covariance(z_mean, z_cov, A, b) return gv.GaussianVar(y_mean, y_cov)
def linear_relu(x, A, b): """ compute y = relu(x)^T A + b """ x_var_diag = tf.matrix_diag_part(x.var) sqrt_x_var_diag = tf.sqrt(x_var_diag) mu = x.mean / (sqrt_x_var_diag + EPSILON) def relu_covariance(x): mu1 = tf.expand_dims(mu, 2) mu2 = tf.transpose(mu1, [0, 2, 1]) s11s22 = tf.expand_dims(x_var_diag, axis=2) * tf.expand_dims( x_var_diag, axis=1) rho = x.var / (tf.sqrt(s11s22)) # + EPSILON) rho = tf.clip_by_value(rho, -1 / (1 + EPSILON), 1 / (1 + EPSILON)) return x.var * bu.delta(rho, mu1, mu2) z_mean = sqrt_x_var_diag * bu.softrelu(mu) y_mean = tf.matmul(z_mean, A.mean) + b.mean z_cov = relu_covariance(x) y_cov = linear_covariance(z_mean, z_cov, A, b) return gv.GaussianVar(y_mean, y_cov)
def simple(x, A, b): mu = x.mean y_mean = tf.matmul(mu, A.mean) + b.mean y_cov = x.var return gv.GaussianVar(y_mean, y_cov)