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
