def testContinuousBernoulliContinuousBernoulliKL(self):
batch_size = 6
a_p = np.array([0.6] * batch_size, dtype=np.float32)
b_p = np.array([0.4] * batch_size, dtype=np.float32)
a = tfd.ContinuousBernoulli(probs=a_p, validate_args=True)
b = tfd.ContinuousBernoulli(probs=b_p, validate_args=True)
kl = tfd.kl_divergence(a, b)
kl_val = self.evaluate(kl)
kl_expected = (
mean(a_p)
* (
np.log(a_p)
+ np.log1p(-b_p)
- np.log(b_p)
- np.log1p(-a_p)
)
+ log_norm_const(a_p)
- log_norm_const(b_p)
+ np.log1p(-a_p)
- np.log1p(-b_p)
)
self.assertEqual(kl.shape, (batch_size,))
self.assertAllClose(kl_val, kl_expected)
def testSlice(self):
logits = self.evaluate(
tf.random.normal([20, 3, 1, 5], seed=test_util.test_seed()))
dist = tfd.ContinuousBernoulli(logits=logits, validate_args=True)
batch_shape = tensorshape_util.as_list(dist.batch_shape)
dist_noshape = tfd.ContinuousBernoulli(
logits=tf1.placeholder_with_default(logits, shape=None),
validate_args=True)
def check(*slicers):
for ds, assert_static_shape in (dist, True), (dist_noshape, False):
bs = batch_shape
prob = self.evaluate(dist.prob(0))
for slicer in slicers:
ds = slicer(ds)
bs = slicer(np.zeros(bs)).shape
prob = slicer(prob)
if assert_static_shape or tf.executing_eagerly():
self.assertAllEqual(bs, ds.batch_shape)
else:
self.assertIsNone(tensorshape_util.rank(
ds.batch_shape))
self.assertAllEqual(bs,
self.evaluate(ds.batch_shape_tensor()))
self.assertAllClose(prob, self.evaluate(ds.prob(0)))
check(make_slicer[3])
check(make_slicer[tf.newaxis])
check(make_slicer[3::7])
check(make_slicer[:, :2])
check(make_slicer[tf.newaxis, :, ..., 0, :2])
check(make_slicer[tf.newaxis, :, ..., 3:, tf.newaxis])
check(make_slicer[..., tf.newaxis, 3:, tf.newaxis])
check(make_slicer[..., tf.newaxis, -3:, tf.newaxis])
check(make_slicer[tf.newaxis, :-3, tf.newaxis, ...])
def halfway(x):
if isinstance(x, tfd.ContinuousBernoulli):
return x.batch_shape_tensor()[0] // 2
return x.shape[0] // 2
check(lambda x: x[halfway(x)])
check(lambda x: x[:halfway(x)])
check(lambda x: x[halfway(x):])
check(
make_slicer[:-3, tf.newaxis],
make_slicer[..., 0, :2],
make_slicer[::2],
)
if tf.executing_eagerly():
return
with self.assertRaisesRegexp(
(ValueError, tf.errors.InvalidArgumentError),
'Index out of range.*input has only 4 dims'):
check(make_slicer[19, tf.newaxis, 2, ..., :, 0, 4])
with self.assertRaisesRegexp(
(ValueError, tf.errors.InvalidArgumentError),
'slice index.*out of bounds',
):
check(make_slicer[..., 2, :]) # ...,1,5 -> 2 is oob.
def testLogits(self):
logits = [-42.0, 42.0]
dist = tfd.ContinuousBernoulli(logits=logits, validate_args=True)
self.assertAllClose(logits, self.evaluate(dist.logits))
self.assertAllClose(sp_special.expit(logits),
self.evaluate(dist.probs_parameter()))
prob = [0.01, 0.99, 0.42]
dist = tfd.ContinuousBernoulli(probs=prob, validate_args=True)
self.assertAllClose(sp_special.logit(prob),
self.evaluate(dist.logits_parameter()))
def testQuantileAtExtremesIsNotNaN(self):
prob = [[0.], [1.]]
dist = tfd.ContinuousBernoulli(probs=prob, validate_args=True)
self.assertAllNotNan(
self.evaluate(
dist.quantile(
np.array([[0., 0.1, 0.3, 0.9, 1.]], dtype=np.float32))))
def testDocstrSliceExample(self):
# batch shape [3, 5, 7, 9]
cb = tfd.ContinuousBernoulli(logits=tf.zeros([3, 5, 7, 9]),
validate_args=True)
self.assertAllEqual((3, 5, 7, 9), cb.batch_shape)
cb2 = cb[:, tf.newaxis, ..., -2:, 1::2] # batch shape [3, 1, 5, 2, 4]
self.assertAllEqual((3, 1, 5, 2, 4), cb2.batch_shape)
def _testPdf(self, **kwargs):
dist = tfd.ContinuousBernoulli(validate_args=True, **kwargs)
# pylint: disable=bad-continuation
xs = [
0.1,
[0.9],
[0.9, 0.1],
[[0.9, 0.1]],
[[0.9, 0.1], [0.9, 0.9]],
]
expected_pmfs = [
[[pdf(0.1, 0.2), pdf(0.1, 0.4)], [pdf(0.1, 0.3),
pdf(0.1, 0.6)]],
[[pdf(0.9, 0.2), pdf(0.9, 0.4)], [pdf(0.9, 0.3),
pdf(0.9, 0.6)]],
[[pdf(0.9, 0.2), pdf(0.1, 0.4)], [pdf(0.9, 0.3),
pdf(0.1, 0.6)]],
[[pdf(0.9, 0.2), pdf(0.1, 0.4)], [pdf(0.9, 0.3),
pdf(0.1, 0.6)]],
[[pdf(0.9, 0.2), pdf(0.1, 0.4)], [pdf(0.9, 0.3),
pdf(0.9, 0.6)]],
]
# pylint: enable=bad-continuation
for x, expected_pmf in zip(xs, expected_pmfs):
self.assertAllClose(self.evaluate(dist.prob(x)), expected_pmf)
self.assertAllClose(self.evaluate(dist.log_prob(x)),
np.log(expected_pmf))
def testCdfMultidimensional(self):
probs = np.array([0., 0.1, 0.3, 0.7, 0.9, 1.],
dtype=np.float32)[..., np.newaxis]
logits = np.log(probs) - np.log1p(-probs)
dist = tfd.ContinuousBernoulli(logits=logits)
x = np.array([-1., 0., 0.1, 0.3, 0.7, 0.9, 1., 2.], dtype=np.float32)
self.assertAllClose(self.evaluate(dist.cdf(x)), cdf(x, probs))
def testLogNormalizerIsNotInf(self):
# Parameterize by logits large in magnitude.
logits = np.array([-1000., -100., 100., 1000.], dtype=np.float32)
dist = tfd.ContinuousBernoulli(logits=logits)
log_normalizer_ = self.evaluate(dist._log_normalizer())
self.assertFalse(np.any(np.isinf(log_normalizer_)))
# The normalizer grows roughly log(abs(logits))
self.assertAllClose(log_normalizer_, np.log(np.abs(logits)))
def testLogNormalizerNearHalfStable(self):
# Parameterize by probs near a half
probs = tf.random.uniform(
[int(1e4)], minval=0.4, maxval=0.6, seed=test_util.test_seed())
dist = tfd.ContinuousBernoulli(probs=probs)
log_normalizer_ = self.evaluate(dist._log_normalizer())
self.assertFalse(np.any(np.isinf(log_normalizer_)))
self.assertFalse(np.any(np.isnan(log_normalizer_)))
def testScalarSlice(self):
logits = self.evaluate(tf.random.normal([], seed=test_util.test_seed()))
dist = tfd.ContinuousBernoulli(logits=logits, validate_args=True)
self.assertAllEqual([], dist.batch_shape)
self.assertAllEqual([1], dist[tf.newaxis].batch_shape)
self.assertAllEqual([], dist[...].batch_shape)
self.assertAllEqual(
[1, 1], dist[tf.newaxis, ..., tf.newaxis].batch_shape)
def testMeanNearHalfStable(self):
# Parameterize by probs near a half
prob = tf.random.uniform(
[int(1e4)], minval=0.4, maxval=0.6, seed=test_util.test_seed())
dist = tfd.ContinuousBernoulli(probs=prob, validate_args=True)
mean_ = self.evaluate(dist.mean())
self.assertFalse(np.any(np.isinf(mean_)))
self.assertFalse(np.any(np.isnan(mean_)))
def testGradientProbs(self):
x = tf.Variable([0.1, 0.7])
self.evaluate(x.initializer)
d = tfd.ContinuousBernoulli(probs=x, validate_args=True)
with tf.GradientTape() as tape:
loss = -d.log_prob([0.1, 0.9])
g = tape.gradient(loss, d.trainable_variables)
self.assertLen(g, 1)
self.assertAllNotNone(g)
def testInvalidProb(self):
invalid_probs = [1.01, 2.0]
for prob in invalid_probs:
with self.assertRaisesOpError(
'probs has components greater than 1'):
dist = tfd.ContinuousBernoulli(probs=prob, validate_args=True)
self.evaluate(dist.probs_parameter())
invalid_probs = [-0.01, -3.0]
for prob in invalid_probs:
with self.assertRaisesOpError('probs has components less than 0'):
dist = tfd.ContinuousBernoulli(probs=prob, validate_args=True)
self.evaluate(dist.probs_parameter())
valid_probs = [0.1, 0.5, 0.9]
for prob in valid_probs:
dist = tfd.ContinuousBernoulli(probs=prob, validate_args=True)
self.assertEqual(np.float32(prob), self.evaluate(dist.probs))
def testPdfWithFloatArg(self):
prob = [[0.2], [0.4], [0.3], [0.6]]
samps = [0.1, 0.2, 0.8]
self.assertAllClose(
np.log(pdf(samps, prob)),
self.evaluate(
tfd.ContinuousBernoulli(probs=prob,
validate_args=False).log_prob(samps)),
)
def testReparameterized(self):
prob = tf.constant([0.2, 0.6])
_, grad_prob = tfp.math.value_and_gradient(
lambda x: tfd.ContinuousBernoulli(probs=x, validate_args=True).
sample( # pylint: disable=g-long-lambda
100,
seed=test_util.test_seed()),
prob)
self.assertIsNotNone(grad_prob)
def testCopyUnknownRank(self):
logits = tf1.placeholder_with_default(
tf.random.normal([20, 3, 1, 5], seed=test_util.test_seed()),
shape=None)
dist = tfd.ContinuousBernoulli(
logits=logits, name='cb1', validate_args=True)
self.assertIn('cb1', dist.name)
dist = dist.copy(name='cb2')
self.assertIn('cb2', dist.name)
def testPdfShapes(self):
probs = lambda prob: tf1.placeholder_with_default(prob, shape=None)
def dist(p):
return tfd.ContinuousBernoulli(probs=probs(p), validate_args=True)
self.assertEqual(
2, len(self.evaluate(dist([[0.5], [0.5]]).log_prob(0.9)).shape))
dist = tfd.ContinuousBernoulli(probs=0.5, validate_args=True)
self.assertEqual(
2, len(self.evaluate(dist.log_prob([[0.9], [0.9]])).shape))
dist = tfd.ContinuousBernoulli(probs=0.5, validate_args=True)
self.assertAllEqual([], dist.log_prob(0.9).shape)
self.assertAllEqual([1], dist.log_prob([0.9]).shape)
self.assertAllEqual([2, 1], dist.log_prob([[0.9], [0.9]]).shape)
dist = tfd.ContinuousBernoulli(probs=[[0.5], [0.5]], validate_args=True)
self.assertAllEqual([2, 1], dist.log_prob(0.9).shape)
def testQuantile(self):
prob = 0.2
dist = tfd.ContinuousBernoulli(probs=prob, validate_args=True)
self.assertAllClose(
self.evaluate(
dist.quantile(np.array([0.1, 0.3, 0.9], dtype=np.float32))
),
np.array(
[quantile(0.1, 0.2), quantile(0.3, 0.2), quantile(0.9, 0.2)],
dtype=np.float32))
def testSampleN(self):
prob = [0.2, 0.6]
dist = tfd.ContinuousBernoulli(probs=prob, validate_args=True)
n = 100000
samples = dist.sample(n, seed=test_util.test_seed())
tensorshape_util.set_shape(samples, [n, 2])
self.assertEqual(samples.dtype, tf.float32)
sample_values = self.evaluate(samples)
self.assertTrue(np.all(sample_values > 0))
self.assertTrue(np.all(sample_values < 1))
# This tolerance is very sensitive to the value of prob as well as n.
self.assertAllClose(mean(np.array(prob, dtype=np.float32)),
np.mean(sample_values, axis=0),
atol=1e-2)
# In this test we're just interested in verifying there isn't a crash
# owing to mismatched types. b/30940152
dist = tfd.ContinuousBernoulli(np.log([0.2, 0.4]), validate_args=True)
x = dist.sample(1, seed=test_util.test_seed())
self.assertAllEqual((1, 2), tensorshape_util.as_list(x.shape))
def testPdfCorrectBroadcastDynamicShape(self):
prob = tf1.placeholder_with_default([0.2, 0.3, 0.4], shape=None)
dist = tfd.ContinuousBernoulli(probs=prob, validate_args=True)
event1 = [0.9, 0.1, 1.0]
event2 = [[0.9, 0.1, 1.0]]
self.assertAllClose(
[pdf(0.9, 0.2), pdf(0.1, 0.3), pdf(1.0, 0.4)],
self.evaluate(dist.prob(event1)))
self.assertAllClose(
[[pdf(0.9, 0.2), pdf(0.1, 0.3), pdf(1.0, 0.4)]],
self.evaluate(dist.prob(event2)))
def testLogitsNearZeroNotNaN(self):
dist = tfd.ContinuousBernoulli(logits=[-0.1, 0., 0.1])
self.assertAllNotNan(self.evaluate(dist.mean()))
self.assertAllNotNan(self.evaluate(dist.variance()))
samples = self.evaluate(dist.sample(int(1e3), seed=test_util.test_seed()))
self.assertAllNotNan(samples)
self.assertAllNotNan(dist.log_prob(samples))
self.assertAllNotNan(dist.log_cdf(samples))
# Check specific values.
self.assertAllClose(0.5, dist.mean()[1])
self.assertAllClose(1 / 12., dist.variance()[1])
def testSliceSequencePreservesOrigVarGradLinkage(self):
logits = tf.Variable(
tf.random.normal([20, 3, 1, 5], seed=test_util.test_seed()))
self.evaluate(logits.initializer)
dist = tfd.ContinuousBernoulli(logits=logits, validate_args=True)
for slicer in [make_slicer[:5], make_slicer[..., -1], make_slicer[:, 1::2]]:
with tf.GradientTape() as tape:
dist = slicer(dist)
lp = dist.log_prob(0.1)
dlpdlogits = tape.gradient(lp, logits)
self.assertIsNotNone(dlpdlogits)
self.assertGreater(
self.evaluate(tf.reduce_sum(tf.abs(dlpdlogits))), 0)
def testSliceCopyOverrideNameSliceAgainCopyOverrideLogitsSliceAgain(self):
seed_stream = test_util.test_seed_stream('slice_continuous_bernoulli')
logits = tf.random.normal([20, 3, 2, 5], seed=seed_stream())
dist = tfd.ContinuousBernoulli(
logits=logits, name='cb1', validate_args=True)
self.assertIn('cb1', dist.name)
dist = dist[:10].copy(name='cb2')
self.assertAllEqual((10, 3, 2, 5), dist.batch_shape)
self.assertIn('cb2', dist.name)
dist = dist.copy(name='cb3')[..., 1]
self.assertAllEqual((10, 3, 2), dist.batch_shape)
self.assertIn('cb3', dist.name)
dist = dist.copy(logits=tf.random.normal([2], seed=seed_stream()))
self.assertAllEqual((2,), dist.batch_shape)
self.assertIn('cb3', dist.name)
def testVarianceAndStd(self):
prob = [[0.2, 0.7], [0.8, 0.4]]
dist = tfd.ContinuousBernoulli(probs=prob, validate_args=True)
self.assertAllClose(
self.evaluate(dist.variance()),
np.array(
[[var(0.2), var(0.7)], [var(0.8), var(0.4)]], dtype=np.float32
), rtol=1e-5)
self.assertAllClose(
self.evaluate(dist.stddev()),
np.array(
[
[np.sqrt(var(0.2)), np.sqrt(var(0.7))],
[np.sqrt(var(0.8)), np.sqrt(var(0.4))],
],
dtype=np.float32), rtol=1e-5)
def testCdfAndLogCdf(self):
prob = 0.2
dist = tfd.ContinuousBernoulli(probs=prob)
self.assertAllClose(
self.evaluate(
dist.cdf(np.array([-2.0, 0.3, 1.1], dtype=np.float32))),
np.array(
[cdf(-2.0, 0.2), cdf(0.3, 0.2),
cdf(1.1, 0.2)],
dtype=np.float32))
self.assertAllClose(
self.evaluate(
dist.log_cdf(np.array([-2.0, 0.3, 1.1], dtype=np.float32))),
np.array([
np.log(cdf(-2.0, 0.2)),
np.log(cdf(0.3, 0.2)),
np.log(cdf(1.1, 0.2)),
],
dtype=np.float32))
def testSampleDeterministicScalarVsVector(self):
prob = [0.2, 0.6]
dist = tfd.ContinuousBernoulli(probs=prob, validate_args=True)
n = 1000
def _seed(seed=None):
seed = test_util.test_seed() if seed is None else seed
if tf.executing_eagerly():
tf1.set_random_seed(seed)
return seed
seed = _seed()
self.assertAllClose(self.evaluate(dist.sample(n, _seed(seed=seed))),
self.evaluate(dist.sample([n], _seed(seed=seed))))
n = tf1.placeholder_with_default(np.int32(1000), shape=None)
seed = _seed()
sample1 = dist.sample(n, _seed(seed=seed))
sample2 = dist.sample([n], _seed(seed=seed))
sample1, sample2 = self.evaluate([sample1, sample2])
self.assertAllClose(sample1, sample2)
def testSliceThenCopyPreservesOrigVarGradLinkage(self):
logits = tf.Variable(
tf.random.normal([20, 3, 1, 5], seed=test_util.test_seed()))
self.evaluate(logits.initializer)
dist = tfd.ContinuousBernoulli(logits=logits, validate_args=True)
dist = dist[:5]
with tf.GradientTape() as tape:
dist = dist.copy(name='contbern2')
lp = dist.log_prob(0.1)
dlpdlogits = tape.gradient(lp, logits)
self.assertIn('contbern2', dist.name)
self.assertIsNotNone(dlpdlogits)
self.assertGreater(self.evaluate(tf.reduce_sum(tf.abs(dlpdlogits))), 0)
with tf.GradientTape() as tape:
dist = dist[:3]
lp = dist.log_prob(0)
dlpdlogits = tape.gradient(lp, logits)
self.assertIn('contbern2', dist.name)
self.assertIsNotNone(dlpdlogits)
self.assertGreater(self.evaluate(tf.reduce_sum(tf.abs(dlpdlogits))), 0)
def testSurvivalAndLogSurvival(self):
prob = 0.2
dist = tfd.ContinuousBernoulli(probs=prob)
self.assertAllClose(
self.evaluate(
dist.survival_function(
np.array([-2.0, 0.3, 1.1], dtype=np.float32))),
1.0 - np.array(
[cdf(-2.0, 0.2), cdf(0.3, 0.2),
cdf(1.1, 0.2)],
dtype=np.float32))
self.assertAllClose(
self.evaluate(
dist.log_survival_function(
np.array([-2.0, 0.3, 1.1], dtype=np.float32))),
np.array([
np.log(1.0 - cdf(-2.0, 0.2)),
np.log(1.0 - cdf(0.3, 0.2)),
np.log(1.0 - cdf(1.1, 0.2)),
],
dtype=np.float32))
def testProbs(self): prob = [0.2, 0.4] dist = tfd.ContinuousBernoulli(probs=prob, validate_args=True) self.assertAllClose(prob, self.evaluate(dist.probs))
def testMean(self): prob = np.array([[0.2, 0.7], [0.8, 0.4]], dtype=np.float32) dist = tfd.ContinuousBernoulli(probs=prob, validate_args=True) self.assertAllClose(self.evaluate(dist.mean()), mean(prob))
