def test_unknown_event_rank(self):
if tf.executing_eagerly():
self.skipTest('Eager execution.')
unknown_rank_dist = tfd.Independent(
tfd.Normal(loc=tf.ones([2, 1, 3]), scale=2.),
reinterpreted_batch_ndims=tf1.placeholder_with_default(1, shape=[]))
td = tfd.TransformedDistribution(
distribution=unknown_rank_dist,
bijector=tfb.Scale(1.),
validate_args=True)
self.assertEqual(td.batch_shape, tf.TensorShape(None))
self.assertEqual(td.event_shape, tf.TensorShape(None))
self.assertAllEqual(td.batch_shape_tensor(), [2, 1])
self.assertAllEqual(td.event_shape_tensor(), [3])
joint_td = tfd.TransformedDistribution(
distribution=tfd.JointDistributionSequentialAutoBatched(
[unknown_rank_dist, unknown_rank_dist]),
bijector=tfb.Invert(tfb.Split(2)),
validate_args=True)
# Note that the current behavior is conservative; we could also correctly
# return a batch shape of `[]` in this case.
self.assertEqual(joint_td.batch_shape, tf.TensorShape(None))
self.assertEqual(joint_td.event_shape, tf.TensorShape(None))
self.assertAllEqual(joint_td.batch_shape_tensor(), [])
self.assertAllEqual(joint_td.event_shape_tensor(), [2, 1, 6])
def testLogProb(self, event_shape, event_dims, training, layer_cls):
training = tf.compat.v1.placeholder_with_default(
training, (), "training")
layer = layer_cls(axis=event_dims, epsilon=0.)
batch_norm = tfb.BatchNormalization(batchnorm_layer=layer,
training=training)
base_dist = distributions.MultivariateNormalDiag(
loc=np.zeros(np.prod(event_shape), dtype=np.float32))
# Reshape the events.
if isinstance(event_shape, int):
event_shape = [event_shape]
base_dist = distributions.TransformedDistribution(
distribution=base_dist,
bijector=tfb.Reshape(event_shape_out=event_shape))
dist = distributions.TransformedDistribution(distribution=base_dist,
bijector=batch_norm,
validate_args=True)
samples = dist.sample(int(1e5))
# No volume distortion since training=False, bijector is initialized
# to the identity transformation.
base_log_prob = base_dist.log_prob(samples)
dist_log_prob = dist.log_prob(samples)
self.evaluate(tf.compat.v1.global_variables_initializer())
base_log_prob_, dist_log_prob_ = self.evaluate(
[base_log_prob, dist_log_prob])
self.assertAllClose(base_log_prob_, dist_log_prob_)
def testStddev(self):
base_stddev = 2.
shift = np.array([[-1, 0, 1], [-1, -2, -3]], dtype=np.float32)
scale = np.array([[1, -2, 3], [2, -3, 2]], dtype=np.float32)
expected_stddev = tf.abs(base_stddev * scale)
normal = tfd.TransformedDistribution(
distribution=tfd.Normal(loc=tf.zeros_like(shift),
scale=base_stddev * tf.ones_like(scale),
validate_args=True),
bijector=tfb.Chain([tfb.Shift(shift=shift),
tfb.Scale(scale=scale)],
validate_args=True),
validate_args=True)
self.assertAllClose(expected_stddev, normal.stddev())
self.assertAllClose(expected_stddev**2, normal.variance())
split_normal = tfd.TransformedDistribution(
distribution=tfd.Independent(normal, reinterpreted_batch_ndims=1),
bijector=tfb.Split(3),
validate_args=True)
self.assertAllCloseNested(tf.split(expected_stddev,
num_or_size_splits=3,
axis=-1),
split_normal.stddev())
scaled_normal = tfd.TransformedDistribution(
distribution=tfd.Independent(normal, reinterpreted_batch_ndims=1),
bijector=tfb.ScaleMatvecTriL([[1., 0.], [-1., 2.]]),
validate_args=True)
with self.assertRaisesRegex(
NotImplementedError, 'is a multivariate transformation'):
scaled_normal.stddev()
def testCovarianceNotImplemented(self):
mvn = tfd.MultivariateNormalDiag(loc=[0., 0.], scale_diag=[1., 2.])
# Non-affine bijector.
with self.assertRaisesRegex(
NotImplementedError, '`covariance` is not implemented'):
tfd.TransformedDistribution(
distribution=mvn, bijector=tfb.Exp()).covariance()
# Non-injective bijector.
with self.assertRaisesRegex(
NotImplementedError, '`covariance` is not implemented'):
tfd.TransformedDistribution(
distribution=mvn, bijector=tfb.AbsoluteValue()).covariance()
# Non-vector event shape.
with self.assertRaisesRegex(
NotImplementedError, '`covariance` is only implemented'):
tfd.TransformedDistribution(
distribution=mvn,
bijector=tfb.Reshape(event_shape_out=[2, 1],
event_shape_in=[2])).covariance()
# Multipart bijector.
with self.assertRaisesRegex(
NotImplementedError, '`covariance` is only implemented'):
tfd.TransformedDistribution(
distribution=mvn, bijector=tfb.Split(2)).covariance()
def test_transformed_affine(self):
sample_shape = 3
mvn = tfd.Independent(tfd.Normal(loc=[0., 0], scale=1), 1)
aff = tfb.Affine(scale_tril=[[0.75, 0.], [0.05, 0.5]])
def expected_lp(y):
x = aff.inverse(y) # Ie, tf.random.normal([4, 3, 2])
fldj = aff.forward_log_det_jacobian(x, event_ndims=1)
return tf.reduce_sum(mvn.log_prob(x) - fldj, axis=1)
# Transform a Sample.
d = tfd.TransformedDistribution(tfd.Sample(mvn,
sample_shape,
validate_args=True),
bijector=aff)
y = d.sample(4, seed=test_util.test_seed())
actual_lp = d.log_prob(y)
self.assertAllEqual((4, ) + (sample_shape, ) + (2, ), y.shape)
self.assertAllEqual((4, ), actual_lp.shape)
self.assertAllClose(*self.evaluate([expected_lp(y), actual_lp]),
atol=0.,
rtol=1e-3)
# Sample a Transform.
d = tfd.Sample(tfd.TransformedDistribution(mvn, bijector=aff),
sample_shape,
validate_args=True)
y = d.sample(4, seed=test_util.test_seed())
actual_lp = d.log_prob(y)
self.assertAllEqual((4, ) + (sample_shape, ) + (2, ), y.shape)
self.assertAllEqual((4, ), actual_lp.shape)
self.assertAllClose(*self.evaluate([expected_lp(y), actual_lp]),
atol=0.,
rtol=1e-3)
def test_transform_parts_to_vector(self, known_split_sizes):
batch_shape = [4, 2]
true_split_sizes = [1, 3, 2]
# Create a joint distribution with parts of the specified sizes.
seed = test_util.test_seed_stream()
component_dists = tf.nest.map_structure(
lambda size: tfd.MultivariateNormalDiag( # pylint: disable=g-long-lambda
loc=tf.random.normal(batch_shape + [size], seed=seed()),
scale_diag=tf.exp(
tf.random.normal(batch_shape + [size], seed=seed()))),
true_split_sizes)
base_dist = tfd.JointDistributionSequential(component_dists)
# Transform to a vector-valued distribution by concatenating the parts.
bijector = tfb.Invert(tfb.Split(known_split_sizes, axis=-1))
with self.assertRaisesRegexp(ValueError, 'Overriding the batch shape'):
tfd.TransformedDistribution(base_dist, bijector, batch_shape=[3])
with self.assertRaisesRegexp(ValueError, 'Overriding the event shape'):
tfd.TransformedDistribution(base_dist, bijector, event_shape=[3])
concat_dist = tfd.TransformedDistribution(base_dist, bijector)
self.assertAllEqual(concat_dist.event_shape, [sum(true_split_sizes)])
self.assertAllEqual(self.evaluate(concat_dist.event_shape_tensor()),
[sum(true_split_sizes)])
self.assertAllEqual(concat_dist.batch_shape, batch_shape)
self.assertAllEqual(self.evaluate(concat_dist.batch_shape_tensor()),
batch_shape)
# Since the Split bijector has (constant) unit Jacobian, the transformed
# entropy and mean/mode should match the base entropy and (split) base
# mean/mode.
self.assertAllEqual(*self.evaluate(
(base_dist.entropy(), concat_dist.entropy())))
self.assertAllEqual(*self.evaluate(
(concat_dist.mean(), bijector.forward(base_dist.mean()))))
self.assertAllEqual(*self.evaluate(
(concat_dist.mode(), bijector.forward(base_dist.mode()))))
# Since the Split bijector has zero Jacobian, the transformed `log_prob`
# and `prob` should match the base distribution.
sample_shape = [3]
x = base_dist.sample(sample_shape, seed=seed())
y = bijector.forward(x)
for attr in ('log_prob', 'prob'):
base_attr = getattr(base_dist, attr)(x)
concat_attr = getattr(concat_dist, attr)(y)
self.assertAllClose(*self.evaluate((base_attr, concat_attr)))
# Test that `.sample()` works and returns a result of the expected structure
# and shape.
y_sampled = concat_dist.sample(sample_shape, seed=seed())
self.assertAllEqual(y.shape, y_sampled.shape)
def testTransformedKLDifferentBijectorFails(self):
d1 = tfd.TransformedDistribution(
tfd.Exponential(rate=0.25),
bijector=tfb.Scale(scale=2.),
validate_args=True)
d2 = tfd.TransformedDistribution(
tfd.Exponential(rate=0.25),
bijector=tfb.Scale(scale=3.),
validate_args=True)
with self.assertRaisesRegex(
NotImplementedError, r'their bijectors are not equal'):
tfd.kl_divergence(d1, d2)
def testTransformedDistribution(self):
mu = 3.0
sigma = 2.0
# Note: the Jacobian callable only works for this example; more generally
# you may or may not need a reduce_sum.
log_normal = tfd.TransformedDistribution(
distribution=tfd.Normal(loc=mu, scale=sigma),
bijector=tfb.Exp(),
validate_args=True)
sp_dist = stats.lognorm(s=sigma, scale=np.exp(mu))
# sample
sample = log_normal.sample(100000, seed=test_util.test_seed())
self.assertAllEqual([], log_normal.event_shape)
self.assertAllEqual([], self.evaluate(log_normal.event_shape_tensor()))
self.assertAllClose(
sp_dist.mean(), np.mean(self.evaluate(sample)), atol=0.0, rtol=0.05)
# pdf, log_pdf, cdf, etc...
# The mean of the lognormal is around 148.
test_vals = np.linspace(0.1, 1000., num=20).astype(np.float32)
for func in [[log_normal.log_prob, sp_dist.logpdf],
[log_normal.prob, sp_dist.pdf],
[log_normal.log_cdf, sp_dist.logcdf],
[log_normal.cdf, sp_dist.cdf],
[log_normal.survival_function, sp_dist.sf],
[log_normal.log_survival_function, sp_dist.logsf]]:
actual = func[0](test_vals)
expected = func[1](test_vals)
self.assertAllClose(
expected, self.evaluate(actual), atol=0, rtol=0.01)
def testCachedSamplesInvert(self):
class ExpInverseOnly(tfb.Bijector):
def __init__(self):
parameters = dict(locals())
super(ExpInverseOnly, self).__init__(
inverse_min_event_ndims=0,
parameters=parameters)
def _inverse(self, y):
return tf.math.log(y)
def _inverse_log_det_jacobian(self, y):
return -tf.math.log(y)
exp_inverse_only = ExpInverseOnly()
log_forward_only = tfb.Invert(exp_inverse_only)
# The log bijector isn't defined over the whole real line, so we make
# sigma sufficiently small so that the draws are positive.
mu = 2.
sigma = 1e-2
exp_normal = tfd.TransformedDistribution(
distribution=tfd.Normal(loc=mu, scale=sigma),
bijector=log_forward_only,
validate_args=True)
sample = exp_normal.sample(
[2, 3], seed=test_util.test_seed(hardcoded_seed=42))
sample_val, log_pdf_val = self.evaluate(
[sample, exp_normal.log_prob(sample)])
expected_log_pdf = sample_val + stats.norm.logpdf(
np.exp(sample_val), loc=mu, scale=sigma)
self.assertAllClose(expected_log_pdf, log_pdf_val, atol=0., rtol=1e-5)
def testSampleAndLogprob(self):
class ExpForwardOnly(tfb.Bijector):
def __init__(self):
super(ExpForwardOnly, self).__init__(forward_min_event_ndims=0)
def _forward(self, x):
return tf.exp(x)
def _forward_log_det_jacobian(self, x):
return tf.convert_to_tensor(value=x)
exp_forward_only = ExpForwardOnly()
mu = 3.0
sigma = 0.02
log_normal = tfd.TransformedDistribution(
distribution=tfd.Normal(loc=mu, scale=sigma),
bijector=exp_forward_only)
sample, log_pdf = self.evaluate(log_normal.experimental_sample_and_log_prob(
[2, 3], seed=test_util.test_seed()))
expected_log_pdf = stats.lognorm.logpdf(
sample, s=sigma, scale=np.exp(mu))
self.assertAllClose(expected_log_pdf, log_pdf, rtol=1e-4, atol=0.)
sample, log_pdf = self.evaluate(
log_normal.experimental_sample_and_log_prob(seed=test_util.test_seed()))
expected_log_pdf = stats.lognorm.logpdf(
sample, s=sigma, scale=np.exp(mu))
self.assertAllClose(expected_log_pdf, log_pdf, rtol=1e-4, atol=0.)
sample2 = self.evaluate(log_normal.sample(seed=test_util.test_seed()))
self.assertAllClose(sample, sample2, rtol=1e-4)
def testQuantileDescending(self):
td = tfd.TransformedDistribution(
distribution=tfd.Normal(loc=0., scale=[1., 1.]),
bijector=tfb.Shift(shift=1.)(tfb.Scale(scale=[2., -2.])),
validate_args=True)
self.assertAllEqual(tf.ones(td.batch_shape, dtype=tf.bool),
td.quantile(.8) < td.quantile(.9))
def _build_posterior_for_one_parameter(param, batch_shape, seed):
"""Built a transformed-normal variational dist over a parameter's support."""
# Build a trainable Normal distribution.
initial_loc = sample_uniform_initial_state(param,
init_sample_shape=batch_shape,
return_constrained=False,
seed=seed)
loc = tf.Variable(initial_value=initial_loc, name=param.name + '_loc')
scale = tfp_util.DeferredTensor(
tf.nn.softplus,
tf.Variable(initial_value=tf.fill(tf.shape(initial_loc),
value=tf.constant(
-4, initial_loc.dtype)),
name=param.name + '_scale'))
posterior_dist = tfd.Normal(loc=loc, scale=scale)
# Ensure the `event_shape` of the variational distribution matches the
# parameter.
if (param.prior.event_shape.ndims is None
or param.prior.event_shape.ndims > 0):
posterior_dist = tfd.Independent(
posterior_dist,
reinterpreted_batch_ndims=param.prior.event_shape.ndims)
# Transform to constrained parameter space.
posterior_dist = tfd.TransformedDistribution(posterior_dist,
param.bijector,
name='{}_posterior'.format(
param.name))
return posterior_dist
def testMaximumLikelihoodTraining(self):
# Test Maximum Likelihood training with default bijector.
training = tf.placeholder_with_default(True, (), "training")
base_dist = distributions.MultivariateNormalDiag(loc=[0., 0.])
batch_norm = tfb.BatchNormalization(training=training)
dist = distributions.TransformedDistribution(distribution=base_dist,
bijector=batch_norm)
target_dist = distributions.MultivariateNormalDiag(loc=[1., 2.])
target_samples = target_dist.sample(200)
dist_samples = dist.sample(3000)
loss = -tf.reduce_mean(dist.log_prob(target_samples))
with tf.control_dependencies(batch_norm.batchnorm.updates):
train_op = tf.train.AdamOptimizer(1e-2).minimize(loss)
moving_mean = tf.identity(batch_norm.batchnorm.moving_mean)
moving_var = tf.identity(batch_norm.batchnorm.moving_variance)
self.evaluate(tf.global_variables_initializer())
for _ in range(3000):
self.evaluate(train_op)
[dist_samples_, moving_mean_,
moving_var_] = self.evaluate([dist_samples, moving_mean, moving_var])
self.assertAllClose([1., 2.],
np.mean(dist_samples_, axis=0),
atol=5e-2)
self.assertAllClose([1., 2.], moving_mean_, atol=5e-2)
self.assertAllClose([1., 1.], moving_var_, atol=5e-2)
def testExcessiveConcretizationOfParams(self):
loc = tfp_hps.defer_and_count_usage(
tf.Variable(0., name='loc', dtype=tf.float32, shape=self.shape))
scale = tfp_hps.defer_and_count_usage(
tf.Variable(2., name='scale', dtype=tf.float32, shape=self.shape))
bij_scale = tfp_hps.defer_and_count_usage(
tf.Variable(2., name='bij_scale', dtype=tf.float32, shape=self.shape))
event_shape = tfp_hps.defer_and_count_usage(
tf.Variable([2, 2], name='input_event_shape', dtype=tf.int32,
shape=self.shape))
batch_shape = tfp_hps.defer_and_count_usage(
tf.Variable([4, 3, 5], name='input_batch_shape', dtype=tf.int32,
shape=self.shape))
dist = tfd.TransformedDistribution(
distribution=tfd.Normal(loc=loc, scale=scale, validate_args=True),
bijector=tfb.Scale(scale=bij_scale, validate_args=True),
event_shape=event_shape,
batch_shape=batch_shape,
validate_args=True)
for method in ('mean', 'entropy', 'event_shape_tensor',
'batch_shape_tensor'):
with tfp_hps.assert_no_excessive_var_usage(
method, max_permissible=self.max_permissible[method]):
getattr(dist, method)()
with tfp_hps.assert_no_excessive_var_usage(
'sample', max_permissible=self.max_permissible['sample']):
dist.sample(seed=test_util.test_seed())
for method in ('log_prob', 'prob'):
with tfp_hps.assert_no_excessive_var_usage(
method, max_permissible=self.max_permissible[method]):
getattr(dist, method)(np.ones((4, 3, 5, 2, 2)) / 3.)
def testNonInjectiveTransformedDistribution(self):
mu = 1.
sigma = 2.0
abs_normal = tfd.TransformedDistribution(
distribution=tfd.Normal(loc=mu, scale=sigma),
bijector=tfb.AbsoluteValue(),
validate_args=True)
sp_normal = stats.norm(mu, sigma)
# sample
sample = abs_normal.sample(100000, seed=test_util.test_seed())
self.assertAllEqual([], abs_normal.event_shape)
sample_ = self.evaluate(sample)
self.assertAllEqual([], self.evaluate(abs_normal.event_shape_tensor()))
# Abs > 0, duh!
np.testing.assert_array_less(0, sample_)
# Let X ~ Normal(mu, sigma), Y := |X|, then
# P[Y < 0.77] = P[-0.77 < X < 0.77]
self.assertAllClose(
sp_normal.cdf(0.77) - sp_normal.cdf(-0.77),
(sample_ < 0.77).mean(), rtol=0.01)
# p_Y(y) = p_X(-y) + p_X(y),
self.assertAllClose(
sp_normal.pdf(1.13) + sp_normal.pdf(-1.13),
self.evaluate(abs_normal.prob(1.13)))
# Log[p_Y(y)] = Log[p_X(-y) + p_X(y)]
self.assertAllClose(
np.log(sp_normal.pdf(2.13) + sp_normal.pdf(-2.13)),
self.evaluate(abs_normal.log_prob(2.13)))
def _build_trainable_posterior(param, initial_loc_fn):
"""Built a transformed-normal variational dist over a parameter's support."""
loc = tf.compat.v1.get_variable(
param.name + '_loc',
initializer=lambda: initial_loc_fn(param),
dtype=param.prior.dtype,
use_resource=True)
scale = tf.nn.softplus(
tf.compat.v1.get_variable(
param.name + '_scale',
initializer=lambda: -4 * tf.ones_like(initial_loc_fn(param)),
dtype=param.prior.dtype,
use_resource=True))
q = tfd.Normal(loc=loc, scale=scale)
# Ensure the `event_shape` of the variational distribution matches the
# parameter.
if (param.prior.event_shape.ndims is None
or param.prior.event_shape.ndims > 0):
q = tfd.Independent(
q, reinterpreted_batch_ndims=param.prior.event_shape.ndims)
# Transform to constrained parameter space.
return tfd.TransformedDistribution(q, param.bijector)
def testCompareToBijector(self):
"""Demonstrates equivalence between TD, Bijector approach and AR dist."""
sample_shape = np.int32([4, 5])
batch_shape = np.int32([])
event_size = np.int32(2)
batch_event_shape = np.concatenate([batch_shape, [event_size]], axis=0)
sample0 = tf.zeros(batch_event_shape)
affine = tfb.ScaleMatvecTriL(
scale_tril=self._random_scale_tril(event_size), validate_args=True)
ar = tfd.Autoregressive(self._normal_fn(affine),
sample0,
validate_args=True)
ar_flow = tfb.MaskedAutoregressiveFlow(
is_constant_jacobian=True,
shift_and_log_scale_fn=lambda x: [None, affine.forward(x)],
validate_args=True)
td = tfd.TransformedDistribution(
# TODO(b/137665504): Use batch-adding meta-distribution to set the batch
# shape instead of tf.zeros.
distribution=tfd.Sample(tfd.Normal(tf.zeros(batch_shape), 1.),
[event_size]),
bijector=ar_flow,
validate_args=True)
x_shape = np.concatenate([sample_shape, batch_shape, [event_size]],
axis=0)
x = 2. * self._rng.random_sample(x_shape).astype(np.float32) - 1.
td_log_prob_, ar_log_prob_ = self.evaluate(
[td.log_prob(x), ar.log_prob(x)])
self.assertAllClose(td_log_prob_, ar_log_prob_, atol=0., rtol=1e-6)
def test_doc_string(self):
# Generate data.
n = 2000
x2 = np.random.randn(n).astype(dtype=np.float32) * 2.
x1 = np.random.randn(n).astype(dtype=np.float32) + (x2 * x2 / 4.)
data = np.stack([x1, x2], axis=-1)
# Density estimation with MADE.
made = tfb.AutoregressiveNetwork(params=2, hidden_units=[10, 10])
distribution = tfd.TransformedDistribution(
distribution=tfd.Sample(tfd.Normal(0., 1.), [2]),
bijector=tfb.MaskedAutoregressiveFlow(made))
# Construct and fit model.
x_ = tfkl.Input(shape=(2,), dtype=tf.float32)
log_prob_ = distribution.log_prob(x_)
model = tfk.Model(x_, log_prob_)
model.compile(optimizer=tf1.train.AdamOptimizer(),
loss=lambda _, log_prob: -log_prob)
batch_size = 25
model.fit(x=data,
y=np.zeros((n, 0), dtype=np.float32),
batch_size=batch_size,
epochs=1,
steps_per_epoch=1, # Usually `n // batch_size`.
shuffle=True,
verbose=True)
# Use the fitted distribution.
self.assertAllEqual((3, 1, 2), distribution.sample((3, 1)).shape)
self.assertAllEqual(
(3,), distribution.log_prob(np.ones((3, 2), dtype=np.float32)).shape)
def testExcessiveConcretizationOfParamsBatchShapeOverride(self):
# Test methods that are not implemented if event_shape is overriden.
loc = tfp_hps.defer_and_count_usage(
tf.Variable(0., name='loc', dtype=tf.float32, shape=self.shape))
scale = tfp_hps.defer_and_count_usage(
tf.Variable(2., name='scale', dtype=tf.float32, shape=self.shape))
bij_scale = tfp_hps.defer_and_count_usage(
tf.Variable(2.,
name='bij_scale',
dtype=tf.float32,
shape=self.shape))
batch_shape = tfp_hps.defer_and_count_usage(
tf.Variable([4, 3, 5],
name='input_batch_shape',
dtype=tf.int32,
shape=self.shape))
dist = tfd.TransformedDistribution(
distribution=tfd.Normal(loc=loc, scale=scale, validate_args=True),
bijector=tfb.Scale(scale=bij_scale, validate_args=True),
batch_shape=batch_shape,
validate_args=True)
for method in ('log_cdf', 'cdf', 'survival_function',
'log_survival_function'):
with tfp_hps.assert_no_excessive_var_usage(
method, max_permissible=self.max_permissible[method]):
getattr(dist, method)(np.ones((4, 3, 2)) / 3.)
with tfp_hps.assert_no_excessive_var_usage(
'quantile', max_permissible=self.max_permissible['quantile']):
dist.quantile(.1)
def test_batch_broadcast_vector_to_parts(self):
batch_shape = [4, 2]
true_split_sizes = [1, 3, 2]
base_event_size = sum(true_split_sizes)
# Base dist with no batch shape (will require broadcasting).
base_dist = tfd.MultivariateNormalDiag(
loc=tf.random.normal([base_event_size], seed=test_util.test_seed()),
scale_diag=tf.exp(tf.random.normal([base_event_size],
seed=test_util.test_seed())))
# Bijector with batch shape in one part.
bijector = tfb.Chain([tfb.JointMap([tfb.Identity(),
tfb.Identity(),
tfb.Shift(
tf.ones(batch_shape +
[true_split_sizes[-1]]))]),
tfb.Split(true_split_sizes, axis=-1)])
split_dist = tfd.TransformedDistribution(base_dist, bijector)
self.assertAllEqual(split_dist.batch_shape, batch_shape)
# Because one branch of the split has batch shape, TD should feed batches
# of base samples *into* the split, so the batch shape propagates to all
# branches.
xs = split_dist.sample(seed=test_util.test_seed())
self.assertAllEqualNested(
tf.nest.map_structure(lambda x: x.shape, xs),
[batch_shape + [d] for d in true_split_sizes])
def testDocstringExample(self):
exp_gamma_distribution = (
tfd.TransformedDistribution(
distribution=tfd.Gamma(concentration=1., rate=2.),
bijector=tfb.Invert(tfb.Exp())))
self.assertAllEqual(
[], self.evaluate(tf.shape(exp_gamma_distribution.sample())))
def test_doc_string_2(self):
n = 2000
c = np.r_[np.zeros(n // 2), np.ones(n // 2)]
mean_0, mean_1 = 0, 5
x = np.r_[np.random.randn(n // 2).astype(dtype=np.float32) + mean_0,
np.random.randn(n // 2).astype(dtype=np.float32) + mean_1]
shuffle_idxs = np.arange(n)
np.random.shuffle(shuffle_idxs)
x = x[shuffle_idxs]
c = c[shuffle_idxs]
seed = test_util.test_seed_stream()
# Density estimation with MADE.
made = tfb.AutoregressiveNetwork(
params=2,
hidden_units=[1],
event_shape=(1, ),
kernel_initializer=tfk.initializers.VarianceScaling(0.1,
seed=seed() %
2**31),
conditional=True,
conditional_event_shape=(1, ))
distribution = tfd.TransformedDistribution(
distribution=tfd.Sample(tfd.Normal(loc=0., scale=1.),
sample_shape=[1]),
bijector=tfb.MaskedAutoregressiveFlow(made))
# Construct and fit model.
x_ = tfkl.Input(shape=(1, ), dtype=tf.float32)
c_ = tfkl.Input(shape=(1, ), dtype=tf.float32)
log_prob_ = distribution.log_prob(
x_, bijector_kwargs={"conditional_input": c_})
model = tfk.Model([x_, c_], log_prob_)
model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.1),
loss=lambda _, log_prob: -log_prob)
batch_size = 25
model.fit(x=[x, c],
y=np.zeros((n, 0), dtype=np.float32),
batch_size=batch_size,
epochs=3,
steps_per_epoch=n // batch_size,
shuffle=False,
verbose=True)
# Use the fitted distribution to sample condition on c = 1
n_samples = 1000
cond = 1
samples = distribution.sample((n_samples, ),
bijector_kwargs={
"conditional_input":
cond * np.ones((n_samples, 1))
},
seed=seed())
# Assert mean is close to conditional mean
self.assertAllMeansClose(samples[..., 0], mean_1, axis=0, atol=1.)
def disabled_testFailureCase(self):
# TODO(b/140229057): This test should pass.
dist = tfd.Chi(df=np.float32(27.744131))
dist = tfd.TransformedDistribution(
bijector=tfb.NormalCDF(), distribution=dist, batch_shape=[4])
dist = tfb.Expm1()(dist)
samps = 1.7182817 + tf.zeros_like(dist.sample(seed=test_util.test_seed()))
self.assertAllClose(dist.log_prob(samps)[0], dist[0].log_prob(samps[0]))
def disabled_testFailureCase(self): # pylint: disable=invalid-name
# TODO(b/140229057): This test should pass.
dist = tfd.Chi(df=np.float32(27.744131) * np.ones((4,)).astype(np.float32))
dist = tfd.TransformedDistribution(
bijector=tfb.NormalCDF(), distribution=dist)
dist = tfb.Expm1()(dist)
samps = 1.7182817 + tf.zeros_like(dist.sample(seed=test_util.test_seed()))
self.assertAllClose(dist.log_prob(samps)[0], dist[0].log_prob(samps[0]))
def testSupportBijectorOutsideRange(self):
log_normal = tfd.TransformedDistribution(
distribution=tfd.Normal(loc=1., scale=2.),
bijector=tfb.Exp(),
validate_args=True)
x = np.array([-4.2, -1e-6, -1.3])
bijector_inverse_x = (
log_normal._experimental_default_event_space_bijector().inverse(x))
self.assertAllNan(self.evaluate(bijector_inverse_x))
def testConditioning(self):
conditional_normal = tfd.TransformedDistribution(
distribution=tfd.Normal(loc=0., scale=1.),
bijector=_ChooseLocation(loc=[-100., 100.]))
z = [-1, +1, -1, -1, +1]
self.assertAllClose(
np.sign(
self.evaluate(
conditional_normal.sample(5, bijector_kwargs={'z': z}))), z)
def testScalarBatchScalarEventIdentityScale(self):
exp2 = tfd.TransformedDistribution(
tfd.Exponential(rate=0.25),
bijector=tfb.Scale(scale=2.),
validate_args=True)
log_prob = exp2.log_prob(1.)
log_prob_ = self.evaluate(log_prob)
base_log_prob = -0.5 * 0.25 + np.log(0.25)
ildj = np.log(2.)
self.assertAllClose(base_log_prob - ildj, log_prob_, rtol=1e-6, atol=0.)
def testComposeFromTransformedDistribution(self):
actual_log_normal = tfb.Exp()(tfd.TransformedDistribution(
distribution=tfd.Normal(0, 1),
bijector=tfb.AffineScalar(shift=0.5, scale=2.)))
expected_log_normal = tfd.LogNormal(0.5, 2.)
x = tf.constant([0.1, 1., 5.])
self.assertAllClose(
*self.evaluate([actual_log_normal.log_prob(x),
expected_log_normal.log_prob(x)]),
atol=0, rtol=1e-3)
def testQuantile(self):
logit_normal = tfd.TransformedDistribution(
distribution=tfd.Normal(loc=0., scale=1.),
bijector=tfb.Sigmoid(),
validate_args=True)
grid = [0., 0.25, 0.5, 0.75, 1.]
q = logit_normal.quantile(grid)
cdf = logit_normal.cdf(q)
cdf_ = self.evaluate(cdf)
self.assertAllClose(grid, cdf_, rtol=1e-6, atol=0.)
def test_transform_vector_to_parts(self, known_split_sizes):
batch_shape = [4, 2]
true_split_sizes = [1, 3, 2]
base_event_size = sum(true_split_sizes)
base_dist = tfd.MultivariateNormalDiag(
loc=tf.random.normal(batch_shape + [base_event_size],
seed=test_util.test_seed()),
scale_diag=tf.exp(
tf.random.normal(batch_shape + [base_event_size],
seed=test_util.test_seed())))
bijector = tfb.Split(known_split_sizes, axis=-1)
split_dist = tfd.TransformedDistribution(base_dist, bijector)
self.assertRegex(
str(split_dist),
'{}.*batch_shape.*event_shape.*dtype'.format(split_dist.name))
expected_event_shape = [np.array([s]) for s in true_split_sizes]
output_event_shape = [np.array(s) for s in split_dist.event_shape]
self.assertAllEqual(output_event_shape, expected_event_shape)
self.assertAllEqual(self.evaluate(split_dist.event_shape_tensor()),
expected_event_shape)
self.assertAllEqual(split_dist.batch_shape, batch_shape)
self.assertAllEqual(self.evaluate(split_dist.batch_shape_tensor()),
batch_shape)
# Since the Split bijector has (constant) unit Jacobian, the transformed
# entropy and mean/mode should match the base entropy and (split) base
# mean/mode.
self.assertAllEqual(*self.evaluate((base_dist.entropy(),
split_dist.entropy())))
self.assertAllEqualNested(
*self.evaluate((split_dist.mean(),
bijector.forward(base_dist.mean()))))
self.assertAllEqualNested(
*self.evaluate((split_dist.mode(),
bijector.forward(base_dist.mode()))))
# Since the Split bijector has zero Jacobian, the transformed `log_prob`
# and `prob` should match the base distribution.
sample_shape = [3]
x = base_dist.sample(sample_shape, seed=test_util.test_seed())
y = bijector.forward(x)
for attr in ('log_prob', 'prob'):
split_attr = getattr(split_dist, attr)(y)
base_attr = getattr(base_dist, attr)(x)
self.assertAllClose(*self.evaluate((base_attr, split_attr)),
rtol=1e-5)
# Test that `.sample()` works and returns a result of the expected structure
# and shape.
y_sampled = split_dist.sample(sample_shape, seed=test_util.test_seed())
self.assertAllEqual([x.shape for x in y], [x.shape for x in y_sampled])
