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
