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))