def testKLBatchBroadcast(self):
batch_shape = [2]
event_shape = [3]
loc_a, scale_a = self._random_loc_and_scale(batch_shape, event_shape)
# No batch shape.
loc_b, scale_b = self._random_loc_and_scale([], event_shape)
mvn_a = tfd.MultivariateNormalLinearOperator(
loc=loc_a, scale=scale_a, validate_args=True)
mvn_b = tfd.MultivariateNormalLinearOperator(
loc=loc_b, scale=scale_b, validate_args=True)
kl = tfd.kl_divergence(mvn_a, mvn_b)
self.assertEqual(batch_shape, kl.shape)
kl_v = self.evaluate(kl)
expected_kl_0 = self._compute_non_batch_kl(
loc_a[0, :],
self.evaluate(scale_a.to_dense())[0, :, :], loc_b,
self.evaluate(scale_b.to_dense()))
expected_kl_1 = self._compute_non_batch_kl(
loc_a[1, :],
self.evaluate(scale_a.to_dense())[1, :, :], loc_b,
self.evaluate(scale_b.to_dense()))
self.assertAllClose(expected_kl_0, kl_v[0])
self.assertAllClose(expected_kl_1, kl_v[1])
def testNamePropertyIsSetByInitArg(self):
loc = [1., 2.]
scale = tf.linalg.LinearOperatorIdentity(2)
mvn = tfd.MultivariateNormalLinearOperator(loc,
scale,
name='Billy',
validate_args=True)
self.assertStartsWith(mvn.name, 'Billy')
def testVariableLocation(self):
loc = tf.Variable([1., 1.])
scale = tf.linalg.LinearOperatorLowerTriangular(
tf.eye(2), is_non_singular=True)
d = tfd.MultivariateNormalLinearOperator(loc, scale, validate_args=True)
self.evaluate(loc.initializer)
with tf.GradientTape() as tape:
lp = d.log_prob([0., 0.])
self.assertIsNotNone(tape.gradient(lp, loc))
def testMeanAndCovariance(self):
loc, scale = self._random_loc_and_scale(
batch_shape=[3, 4], event_shape=[5])
mvn = tfd.MultivariateNormalLinearOperator(loc, scale, validate_args=True)
self.assertAllEqual(self.evaluate(mvn.mean()), loc)
self.assertAllClose(
self.evaluate(mvn.covariance()),
self.evaluate(scale.matmul(scale, adjoint_arg=True).to_dense()))
def testMeanAndCovariance(self):
loc, scale = self._random_loc_and_scale(
batch_shape=[3, 4], event_shape=[5])
mvn = tfd.MultivariateNormalLinearOperator(loc, scale)
self.assertAllEqual(self.evaluate(mvn.mean()), loc)
self.assertAllClose(
self.evaluate(mvn.covariance()),
np.matmul(
self.evaluate(scale.to_dense()),
np.transpose(self.evaluate(scale.to_dense()), [0, 1, 3, 2])))
def testShapes(self):
loc = self.rng.rand(3, 5, 2)
scale = tf.linalg.LinearOperatorLowerTriangular(
self._random_tril_matrix([3, 5, 2, 2]), is_non_singular=True)
mvn = tfd.MultivariateNormalLinearOperator(loc, scale, validate_args=True)
# Shapes known at graph construction time.
self.assertEqual((2,), tuple(tensorshape_util.as_list(mvn.event_shape)))
self.assertEqual((3, 5), tuple(tensorshape_util.as_list(mvn.batch_shape)))
# Shapes known at runtime.
self.assertEqual((2,), tuple(self.evaluate(mvn.event_shape_tensor())))
self.assertEqual((3, 5), tuple(self.evaluate(mvn.batch_shape_tensor())))
def testVariableScaleAssertions(self):
# We test that changing the scale to be non-invertible raises an exception
# when validate_args is True. This is really just testing the underlying
# AffineLinearOperator instance, but we include it to demonstrate that it
# works as expected.
loc = tf.constant([1., 1.])
scale_tensor = tf.Variable(np.eye(2, dtype=np.float32))
scale = tf.linalg.LinearOperatorLowerTriangular(
scale_tensor,
is_non_singular=True)
d = tfd.MultivariateNormalLinearOperator(loc, scale, validate_args=True)
self.evaluate(scale_tensor.initializer)
with self.assertRaises(Exception):
with tf.control_dependencies([scale_tensor.assign([[1., 0.], [1., 0.]])]):
self.evaluate(d.sample())
def testMeanAndCovarianceAndVarAndStddevLocIsNone(self):
loc = None
scale_shape = (2, 4, 4)
scale = tf.linalg.LinearOperatorLowerTriangular(
self._random_tril_matrix(scale_shape), is_non_singular=True)
mvn = tfd.MultivariateNormalLinearOperator(loc, scale, validate_args=True)
self.assertAllEqual(
self.evaluate(mvn.mean()), np.zeros((2, 4), dtype=np.float32))
expected_cov = self.evaluate(
scale.matmul(scale, adjoint_arg=True).to_dense())
self.assertAllClose(self.evaluate(mvn.covariance()), expected_cov)
expected_var = self.evaluate(tf.linalg.diag_part(expected_cov))
self.assertAllClose(self.evaluate(mvn.variance()), expected_var)
self.assertAllClose(self.evaluate(mvn.stddev()), expected_var**0.5)
def testLogPDFScalarBatch(self):
loc = self.rng.rand(2)
scale = tf.linalg.LinearOperatorLowerTriangular(
self._random_tril_matrix([2, 2]), is_non_singular=True)
mvn = tfd.MultivariateNormalLinearOperator(loc, scale, validate_args=True)
x = self.rng.rand(2)
log_pdf = mvn.log_prob(x)
pdf = mvn.prob(x)
covariance = self.evaluate(
tf.matmul(scale.to_dense(), scale.to_dense(), adjoint_b=True))
scipy_mvn = stats.multivariate_normal(mean=loc, cov=covariance)
expected_log_pdf = scipy_mvn.logpdf(x)
expected_pdf = scipy_mvn.pdf(x)
self.assertEqual((), log_pdf.shape)
self.assertEqual((), pdf.shape)
self.assertAllClose(expected_log_pdf, self.evaluate(log_pdf))
self.assertAllClose(expected_pdf, self.evaluate(pdf))
def testStatsBroadcastWhenLocIsBiggerThanScale(self):
batch_shape = (3,)
event_shape = (4,)
loc_shape = batch_shape + event_shape
loc = self.rng.randn(*loc_shape)
scale_shape = event_shape * 2
scale = tf.linalg.LinearOperatorLowerTriangular(
self._random_tril_matrix(scale_shape),
is_non_singular=True)
mvn = tfd.MultivariateNormalLinearOperator(loc, scale, validate_args=True)
self.assertAllEqual(self.evaluate(mvn.mean()), loc)
cov = mvn.covariance()
self.assertAllEqual(batch_shape + event_shape * 2, cov.shape)
self.assertAllClose(
self.evaluate(cov),
self.evaluate(scale.matmul(scale, adjoint_arg=True).to_dense()) +
np.zeros(loc_shape + (1,), dtype=np.float32))
entropy = mvn.entropy()
self.assertAllEqual(batch_shape, entropy.shape)
def testRaisesIfScaleNotProvided(self):
loc = self.rng.rand(2)
with self.assertRaises(ValueError):
tfd.MultivariateNormalLinearOperator(loc, scale=None, validate_args=True)
