def _apply_variational_kernel(self, inputs): if (not isinstance(self.kernel_posterior, tfd.Independent) or not isinstance(self.kernel_posterior.distribution, tfd.Normal)): raise TypeError( '`DenseFlipout` requires ' '`kernel_posterior_fn` produce an instance of ' '`tfd.Independent(tfd.Normal)` ' '(saw: \"{}\").'.format(self.kernel_posterior.name)) self.kernel_posterior_affine = tfd.Normal( loc=tf.zeros_like(self.kernel_posterior.distribution.loc), scale=self.kernel_posterior.distribution.scale) self.kernel_posterior_affine_tensor = ( self.kernel_posterior_tensor_fn(self.kernel_posterior_affine)) self.kernel_posterior_tensor = None input_shape = tf.shape(inputs) batch_shape = input_shape[:-1] seed_stream = tfd.SeedStream(self.seed, salt='DenseFlipout') sign_input = random_rademacher( input_shape, dtype=inputs.dtype, seed=seed_stream()) sign_output = random_rademacher( tf.concat([batch_shape, tf.expand_dims(self.units, 0)], 0), dtype=inputs.dtype, seed=seed_stream()) perturbed_inputs = self._matmul( inputs * sign_input, self.kernel_posterior_affine_tensor) * sign_output outputs = self._matmul(inputs, self.kernel_posterior.distribution.loc) outputs += perturbed_inputs return outputs
def _apply_variational_kernel(self, inputs): if (not isinstance(self.kernel_posterior, independent_lib.Independent) or not isinstance(self.kernel_posterior.distribution, normal_lib.Normal)): raise TypeError('`DenseFlipout` requires ' '`kernel_posterior_fn` produce an instance of ' '`tfd.Independent(tfd.Normal)` ' '(saw: \"{}\").'.format( self.kernel_posterior.name)) self.kernel_posterior_affine = normal_lib.Normal( loc=tf.zeros_like(self.kernel_posterior.distribution.loc), scale=self.kernel_posterior.distribution.scale) self.kernel_posterior_affine_tensor = (self.kernel_posterior_tensor_fn( self.kernel_posterior_affine)) self.kernel_posterior_tensor = None input_shape = tf.shape(inputs) batch_shape = input_shape[:-1] seed_stream = SeedStream(self.seed, salt='DenseFlipout') sign_input = random_rademacher(input_shape, dtype=inputs.dtype, seed=seed_stream()) sign_output = random_rademacher(tf.concat( [batch_shape, tf.expand_dims(self.units, 0)], 0), dtype=inputs.dtype, seed=seed_stream()) perturbed_inputs = tf.matmul( inputs * sign_input, self.kernel_posterior_affine_tensor) * sign_output outputs = tf.matmul(inputs, self.kernel_posterior.distribution.loc) outputs += perturbed_inputs return outputs
def _apply_variational_kernel(self, inputs): if (not isinstance(self.kernel_posterior, tfd.Independent) or not isinstance(self.kernel_posterior.distribution, tfd.Normal)): raise TypeError(f'`DenseFlipout` requires kernel_posterior_fn` produce an instance of ' f'`tf.distributions.Independent(tf.distributions.Normal)' f'`(saw: \"{self.kernel_posterior.name}\").') self.kernel_posterior_affine = tfd.Normal( loc=tf.zeros_like(self.kernel_posterior.distribution.loc), scale=self.kernel_posterior.distribution.scale) self.kernel_posterior_affine_tensor = (self.kernel_posterior_tensor_fn(self.kernel_posterior_affine)) self.kernel_posterior_tensor = None input_shape = tf.shape(inputs) batch_shape = input_shape[:-1] sign_input = random_rademacher(input_shape, dtype=inputs.dtype) sign_output = random_rademacher( tf.concat([batch_shape, tf.expand_dims(self.units, 0)], 0), dtype=inputs.dtype) perturbed_inputs = self._ndegmul(inputs * sign_input, self.kernel_posterior_affine_tensor) * sign_output outputs = self._ndegmul(inputs, self.kernel_posterior.distribution.loc) outputs += perturbed_inputs return outputs
def _choose_direction_batched(point, seed_stream): with tf.compat.v1.name_scope("choose_direction_batched"): batch_size = tf.shape(input=point.state[0])[0] dtype = point.state[0].dtype return tfp_math.random_rademacher([batch_size], dtype=dtype, seed=seed_stream())
def proposal(seed): """Proposal for log-concave rejection sampler.""" (top_lobe_fractions_seed, exponential_samples_seed, top_selector_seed, random_rademacher_seed) = samplers.split_seed( seed, n=4, salt='log_concave_rejection_sampler_proposal') top_lobe_fractions = samplers.uniform(mode_shape, seed=top_lobe_fractions_seed, dtype=dtype) # V in ref [1]. top_offsets = top_lobe_fractions * top_width / mode_height exponential_samples = exponential_distribution.sample( mode_shape, seed=exponential_samples_seed) # E in ref [1]. exponential_height = ( exponential_distribution.prob(exponential_samples) * mode_height) exponential_offsets = (top_width + exponential_samples) / mode_height top_selector = samplers.uniform(mode_shape, seed=top_selector_seed, dtype=dtype) # U in ref [1]. on_top_mask = tf.less_equal(top_selector, top_fraction) unsigned_offsets = tf.where(on_top_mask, top_offsets, exponential_offsets) offsets = tf.round( tfp_math.random_rademacher( mode_shape, seed=random_rademacher_seed, dtype=dtype) * unsigned_offsets) potential_samples = mode + offsets envelope_height = tf.where(on_top_mask, mode_height, exponential_height) return potential_samples, envelope_height
def _sample_n(self, n, seed=None): # Generate samples using: # mu + sigma* sgn(U-0.5)* sqrt(X^2 + Y^2 + Z^2) U~Unif; X,Y,Z ~N(0,1) normal_seed, rademacher_seed = samplers.split_seed( seed, salt='DoublesidedMaxwell') loc = tf.convert_to_tensor(self.loc) scale = tf.convert_to_tensor(self.scale) shape = prefer_static.pad(self._batch_shape_tensor(loc=loc, scale=scale), paddings=[[1, 0]], constant_values=n) # Generate one-sided Maxwell variables by using 3 Gaussian variates norm_rvs = samplers.normal(shape=prefer_static.pad(shape, paddings=[[0, 1]], constant_values=3), dtype=self.dtype, seed=normal_seed) maxwell_rvs = tf.norm(norm_rvs, axis=-1) # Generate random signs for the symmetric variates. random_sign = tfp_math.random_rademacher(shape, seed=rademacher_seed) sampled = random_sign * maxwell_rvs * scale + loc return sampled
def _uniform_unit_norm(dimension, shape, dtype, seed): """Returns a batch of points chosen uniformly from the unit hypersphere.""" # This works because the Gaussian distribution is spherically symmetric. # raw shape: shape + [dimension] static_dimension = tf.get_static_value(dimension) if static_dimension is not None and static_dimension == 1: return tfp_math.random_rademacher( tf.concat([shape, [1]], axis=0), dtype=dtype, seed=seed) raw = samplers.normal( shape=tf.concat([shape, [dimension]], axis=0), seed=seed, dtype=dtype) unit_norm = raw / tf.norm(raw, ord=2, axis=-1)[..., tf.newaxis] return unit_norm
def call(self, x, sampling=True): """Perform the forward pass""" if sampling: # Flipout-estimated weight samples s = random_rademacher(tf.shape(x), dtype=tf.float64) r = random_rademacher([x.shape[0], self.d_out], dtype=tf.float64) w_samples = tf.nn.softplus(self.w_std) * tf.random.normal( [self.d_in, self.d_out], dtype=tf.float64) w_perturbations = r * tf.matmul(x * s, w_samples) w_outputs = tf.matmul(x, self.w_loc) + w_perturbations # Flipout-estimated bias samples r = random_rademacher([x.shape[0], self.d_out], dtype=tf.float64) b_samples = tf.nn.softplus(self.b_std) * tf.random.normal( [self.d_out], dtype=tf.float64) b_outputs = self.b_loc + r * b_samples return w_outputs + b_outputs else: return x @ self.w_loc + self.b_loc
def test_expected_value(self): shape_ = np.array([2, 3, int(1e3)], np.int32) shape = (tf.constant(shape_) if self.use_static_shape else tf.placeholder_with_default(shape_, shape=None)) x = random_rademacher(shape, self.dtype, seed=42) if self.use_static_shape: self.assertAllEqual(shape_, x.shape) x_ = self.evaluate(x) self.assertEqual(self.dtype, x.dtype.as_numpy_dtype) self.assertAllEqual(shape_, x_.shape) self.assertAllEqual([-1., 1], np.unique(np.reshape(x_, [-1]))) self.assertAllClose(np.zeros(shape_[:-1]), np.mean(x_, axis=-1), atol=0.05, rtol=0.)
def _sample_n(self, n, seed=None, name=None): n = tf.convert_to_tensor(n, name='num', dtype=tf.int32) loc = tf.convert_to_tensor(self.loc) scale = tf.convert_to_tensor(self.scale) power = tf.convert_to_tensor(self.power) batch_shape = self._batch_shape_tensor(loc=loc, scale=scale, power=power) result_shape = prefer_static.concat([[n], batch_shape], axis=0) ipower = tf.broadcast_to(tf.math.reciprocal(power), batch_shape) gamma_dist = gamma.Gamma(ipower, 1.) rademacher_seed, gamma_seed = samplers.split_seed(seed, salt='GenNormal') gamma_sample = gamma_dist.sample(n, seed=gamma_seed) binary_sample = tfp_math.random_rademacher(result_shape, dtype=self.dtype, seed=rademacher_seed) sampled = (binary_sample * tf.math.pow(tf.abs(gamma_sample), ipower)) return loc + scale * sampled