def testPdfUniformZeroBatches(self): # Corresponds to a uniform distribution alpha = [1., 1, 1] x = [[.2, .5, .3], [.3, .4, .3]] dist = dirichlet_lib.Dirichlet(alpha) pdf = dist.prob(x) self.assertAllClose([2., 2.], self.evaluate(pdf)) self.assertEqual((2), pdf.get_shape())
def testEntropy(self):
alpha = [1., 2, 3]
dirichlet = dirichlet_lib.Dirichlet(concentration=alpha)
self.assertEqual(dirichlet.entropy().get_shape(), ())
if not stats:
return
expected_entropy = stats.dirichlet.entropy(alpha)
self.assertAllClose(self.evaluate(dirichlet.entropy()), expected_entropy)
def testModeEnableAllowNanStats(self):
alpha = np.array([1., 2, 3])
dirichlet = dirichlet_lib.Dirichlet(concentration=alpha,
allow_nan_stats=True)
expected_mode = np.zeros_like(alpha) + np.nan
self.assertEqual(dirichlet.mode().get_shape(), [3])
self.assertAllClose(self.evaluate(dirichlet.mode()), expected_mode)
def testPdfZeroBatchesNontrivialX(self):
with self.test_session():
alpha = [1., 2]
x = [.3, .7]
dist = dirichlet_lib.Dirichlet(alpha)
pdf = dist.prob(x)
self.assertAllClose(7. / 5, pdf.eval())
self.assertEqual((), pdf.get_shape())
def testMean(self):
alpha = [1., 2, 3]
dirichlet = dirichlet_lib.Dirichlet(concentration=alpha)
self.assertEqual(dirichlet.mean().get_shape(), [3])
if not stats:
return
expected_mean = stats.dirichlet.mean(alpha)
self.assertAllClose(self.evaluate(dirichlet.mean()), expected_mean)
def testSimpleShapes(self):
with self.test_session():
alpha = np.random.rand(3)
dist = dirichlet_lib.Dirichlet(alpha)
self.assertEqual(3, dist.event_shape_tensor().eval())
self.assertAllEqual([], dist.batch_shape_tensor().eval())
self.assertEqual(tensor_shape.TensorShape([3]), dist.event_shape)
self.assertEqual(tensor_shape.TensorShape([]), dist.batch_shape)
def testPdfZeroBatches(self):
with self.test_session():
alpha = [1., 2]
x = [.5, .5]
dist = dirichlet_lib.Dirichlet(alpha)
pdf = dist.prob(x)
self.assertAllClose(1., pdf.eval())
self.assertEqual((), pdf.get_shape())
def testPdfAlphaStretchedInBroadcastWhenSameRank(self):
with self.test_session():
alpha = [[1., 2]]
x = [[.5, .5], [.3, .7]]
dist = dirichlet_lib.Dirichlet(alpha)
pdf = dist.prob(x)
self.assertAllClose([1., 7. / 5], pdf.eval())
self.assertEqual((2), pdf.get_shape())
def testDirichletFullyReparameterized(self):
alpha = constant_op.constant([1.0, 2.0, 3.0])
with backprop.GradientTape() as tape:
tape.watch(alpha)
dirichlet = dirichlet_lib.Dirichlet(alpha)
samples = dirichlet.sample(100)
grad_alpha = tape.gradient(samples, alpha)
self.assertIsNotNone(grad_alpha)
def testComplexShapes(self):
with self.test_session():
alpha = np.random.rand(3, 2, 2)
dist = dirichlet_lib.Dirichlet(alpha)
self.assertEqual(2, self.evaluate(dist.event_shape_tensor()))
self.assertAllEqual([3, 2],
self.evaluate(dist.batch_shape_tensor()))
self.assertEqual(tensor_shape.TensorShape([2]), dist.event_shape)
self.assertEqual(tensor_shape.TensorShape([3, 2]),
dist.batch_shape)
def testPdfXProper(self):
alpha = [[1., 2, 3]]
dist = dirichlet_lib.Dirichlet(alpha, validate_args=True)
self.evaluate(dist.prob([.1, .3, .6]))
self.evaluate(dist.prob([.2, .3, .5]))
# Either condition can trigger.
with self.assertRaisesOpError("samples must be positive"):
self.evaluate(dist.prob([-1., 1.5, 0.5]))
with self.assertRaisesOpError("samples must be positive"):
self.evaluate(dist.prob([0., .1, .9]))
with self.assertRaisesOpError("sample last-dimension must sum to `1`"):
self.evaluate(dist.prob([.1, .2, .8]))
def testVariance(self):
alpha = [1., 2, 3]
denominator = np.sum(alpha)**2 * (np.sum(alpha) + 1)
dirichlet = dirichlet_lib.Dirichlet(concentration=alpha)
self.assertEqual(dirichlet.covariance().get_shape(), (3, 3))
if not stats:
return
expected_covariance = np.diag(stats.dirichlet.var(alpha))
expected_covariance += [[0., -2, -3], [-2, 0, -6], [-3, -6, 0]
] / denominator
self.assertAllClose(
self.evaluate(dirichlet.covariance()), expected_covariance)
def testVariance(self):
alpha = [1., 2, 3]
denominator = np.sum(alpha)**2 * (np.sum(alpha) + 1)
dirichlet = dirichlet_lib.Dirichlet(concentration=alpha)
# The dirichlet.covariance() will call matmul op on GPU, which might lead to
# deviated results with tf32.
with ops.device('/cpu:0'):
covariance = dirichlet.covariance()
self.assertEqual(covariance.get_shape(), (3, 3))
if not stats:
return
expected_covariance = np.diag(stats.dirichlet.var(alpha))
expected_covariance += [[0., -2, -3], [-2, 0, -6], [-3, -6, 0]
] / denominator
self.assertAllClose(self.evaluate(covariance), expected_covariance)
def testSample(self):
alpha = [1., 2]
dirichlet = dirichlet_lib.Dirichlet(alpha)
n = constant_op.constant(100000)
samples = dirichlet.sample(n)
sample_values = self.evaluate(samples)
self.assertEqual(sample_values.shape, (100000, 2))
self.assertTrue(np.all(sample_values > 0.0))
if not stats:
return
self.assertLess(
stats.kstest(
# Beta is a univariate distribution.
sample_values[:, 0],
stats.beta(a=1., b=2.).cdf)[0],
0.01)
def testCovarianceFromSampling(self):
alpha = np.array([[1., 2, 3], [2.5, 4, 0.01]], dtype=np.float32)
with self.test_session() as sess:
dist = dirichlet_lib.Dirichlet(
alpha) # batch_shape=[2], event_shape=[3]
x = dist.sample(int(250e3), seed=1)
sample_mean = math_ops.reduce_mean(x, 0)
x_centered = x - sample_mean[None, ...]
sample_cov = math_ops.reduce_mean(
math_ops.matmul(x_centered[..., None], x_centered[...,
None, :]), 0)
sample_var = array_ops.matrix_diag_part(sample_cov)
sample_stddev = math_ops.sqrt(sample_var)
[
sample_mean_,
sample_cov_,
sample_var_,
sample_stddev_,
analytic_mean,
analytic_cov,
analytic_var,
analytic_stddev,
] = sess.run([
sample_mean,
sample_cov,
sample_var,
sample_stddev,
dist.mean(),
dist.covariance(),
dist.variance(),
dist.stddev(),
])
self.assertAllClose(sample_mean_,
analytic_mean,
atol=0.,
rtol=0.04)
self.assertAllClose(sample_cov_, analytic_cov, atol=0., rtol=0.06)
self.assertAllClose(sample_var_, analytic_var, atol=0., rtol=0.03)
self.assertAllClose(sample_stddev_,
analytic_stddev,
atol=0.,
rtol=0.02)
def testModeInvalid(self):
alpha = np.array([1., 2, 3])
dirichlet = dirichlet_lib.Dirichlet(
concentration=alpha, allow_nan_stats=False)
with self.assertRaisesOpError("Condition x < y.*"):
self.evaluate(dirichlet.mode())
def testMode(self): alpha = np.array([1.1, 2, 3]) expected_mode = (alpha - 1) / (np.sum(alpha) - 3) dirichlet = dirichlet_lib.Dirichlet(concentration=alpha) self.assertEqual(dirichlet.mode().get_shape(), [3]) self.assertAllClose(self.evaluate(dirichlet.mode()), expected_mode)
def testPdfXStretchedInBroadcastWhenLowerRank(self): alpha = [[1., 2], [2., 3]] x = [.5, .5] pdf = dirichlet_lib.Dirichlet(alpha).prob(x) self.assertAllClose([1., 3. / 2], self.evaluate(pdf)) self.assertEqual((2), pdf.get_shape())
def testConcentrationProperty(self):
alpha = [[1., 2, 3]]
with self.test_session():
dist = dirichlet_lib.Dirichlet(alpha)
self.assertEqual([1, 3], dist.concentration.get_shape())
self.assertAllClose(alpha, self.evaluate(dist.concentration))
def testPdfAlphaStretchedInBroadcastWhenLowerRank(self): alpha = [1., 2] x = [[.5, .5], [.2, .8]] pdf = dirichlet_lib.Dirichlet(alpha).prob(x) self.assertAllClose([1., 8. / 5], self.evaluate(pdf)) self.assertEqual((2), pdf.get_shape())
