def _eval(self, x, weights): kernel, bias = self.unpack_weights_fn(weights) # pylint: disable=not-callable y = x if kernel is not None: kernel_dist, _ = self.unpack_weights_fn( # pylint: disable=not-callable self.posterior.sample_distributions(value=weights)[0]) kernel_loc, kernel_scale = get_spherical_normal_loc_scale( kernel_dist) # batch_size = tf.shape(x)[0] # sign_input_shape = ([batch_size] + # [1] * self._rank + # [self._input_channels]) y *= tfp_random.rademacher(ps.shape(y), dtype=y.dtype, seed=self._seed()) kernel_perturb = normal_lib.Normal(loc=0., scale=kernel_scale) y = self._apply_kernel_fn( # E.g., tf.matmul. y, kernel_perturb.sample(seed=self._seed())) y *= tfp_random.rademacher(ps.shape(y), dtype=y.dtype, seed=self._seed()) y += self._apply_kernel_fn(x, kernel_loc) if bias is not None: y = y + bias if self.activation_fn is not None: y = self.activation_fn(y) # pylint: disable=not-callable return y
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 = tfp_random.rademacher(input_shape, dtype=inputs.dtype, seed=seed_stream()) sign_output = tfp_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 proposal(seed): """Proposal for log-concave rejection sampler.""" (top_lobe_fractions_seed, exponential_samples_seed, top_selector_seed, 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_random.rademacher( mode_shape, seed=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 _choose_direction_batched(point, seed_stream): with tf1.name_scope("choose_direction_batched"): batch_size = tf.shape(input=point.state[0])[0] dtype = point.state[0].dtype return tfp_random.rademacher([batch_size], dtype=dtype, seed=seed_stream())
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_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_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 _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 = ps.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_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