def testLogPDFBatch(self):
seed_stream = test_util.test_seed_stream()
df, loc, scale_row, scale_col = self._random_df_loc_and_scale(
batch_shape=[7], matrix_shape=[3, 5], seed_stream=seed_stream)
matrix_t = tfd.MatrixTLinearOperator(df,
loc,
scale_row,
scale_col,
validate_args=True)
x = tf.random.normal(shape=[3, 5],
seed=seed_stream(),
dtype=self.dtype)
log_pdf = matrix_t.log_prob(x)
pdf = matrix_t.prob(x)
# Check that this matches using a Multivariate-T.
mvt = tfd.MultivariateStudentTLinearOperator(
df,
loc=_vec(loc),
scale=tf.linalg.LinearOperatorKronecker([scale_row, scale_col]))
log_pdf_, pdf_, mvt_log_pdf_, mvt_pdf_ = self.evaluate(
[log_pdf, pdf,
mvt.log_prob(_vec(x)),
mvt.prob(_vec(x))])
self.assertEqual((7, ), log_pdf.shape)
self.assertEqual((7, ), pdf.shape)
self.assertAllClose(mvt_log_pdf_, log_pdf_)
self.assertAllClose(mvt_pdf_, pdf_)
def testShapes(self):
seed_stream = test_util.test_seed_stream()
df = np.array([1., 2., 3.], dtype=self.dtype)
loc = tf.random.normal(shape=[5, 1, 2, 3],
seed=seed_stream(),
dtype=self.dtype)
scale_row = tf.linalg.LinearOperatorLowerTriangular(
self._random_tril_matrix([7, 1, 1, 2, 2], seed_stream()),
is_non_singular=True)
scale_col = tf.linalg.LinearOperatorLowerTriangular(
self._random_tril_matrix([9, 1, 1, 1, 3, 3], seed_stream()),
is_non_singular=True)
matrix_t = tfd.MatrixTLinearOperator(df,
loc,
scale_row,
scale_col,
validate_args=True)
self.assertAllEqual((2, 3), matrix_t.event_shape)
self.assertAllEqual((9, 7, 5, 3), matrix_t.batch_shape)
self.assertAllEqual((2, 3),
self.evaluate(matrix_t.event_shape_tensor()))
self.assertAllEqual((9, 7, 5, 3),
self.evaluate(matrix_t.batch_shape_tensor()))
def testSampleVariance(self):
seed_stream = test_util.test_seed_stream()
df, loc, scale_row, scale_col = self._random_df_loc_and_scale(
batch_shape=[5, 2], matrix_shape=[2, 3], seed_stream=seed_stream)
matrix_t = tfd.MatrixTLinearOperator(df, loc, scale_row, scale_col)
samples = matrix_t.sample(int(2e6), seed=seed_stream())
variance_, samples_ = self.evaluate([matrix_t.variance(), samples])
self.assertAllClose(np.var(samples_, axis=0), variance_, rtol=4e-2)
def testVariableDfAssertions(self):
df = tf.Variable(1.)
loc = tf.constant([[1., 1.]])
scale_row = tf.linalg.LinearOperatorLowerTriangular(
tf.eye(1), is_non_singular=True)
scale_column = tf.linalg.LinearOperatorLowerTriangular(
tf.eye(2), is_non_singular=True)
d = tfd.MatrixTLinearOperator(df,
loc,
scale_row,
scale_column,
validate_args=True)
self.evaluate(df.initializer)
with self.assertRaises(Exception):
with tf.control_dependencies([df.assign(-1.)]):
self.evaluate(d.sample(seed=test_util.test_seed()))
def testVariableLocation(self):
df = tf.constant(1.)
loc = tf.Variable([[1., 1.]])
scale_row = tf.linalg.LinearOperatorLowerTriangular(
tf.eye(1), is_non_singular=True)
scale_column = tf.linalg.LinearOperatorLowerTriangular(
tf.eye(2), is_non_singular=True)
d = tfd.MatrixTLinearOperator(df,
loc,
scale_row,
scale_column,
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 testMeanAndVariance(self):
df, loc, scale_row, scale_col = self._random_df_loc_and_scale(
batch_shape=[3, 4], matrix_shape=[2, 5])
matrix_t = tfd.MatrixTLinearOperator(df, loc, scale_row, scale_col)
cov_row = scale_row.matmul(scale_row.adjoint())
cov_col = scale_col.matmul(scale_col.adjoint())
# Compute diagonal of Kronecker product
expected_variance = (
cov_col.to_dense()[..., :, tf.newaxis, :, tf.newaxis] *
cov_row.to_dense()[..., tf.newaxis, :, tf.newaxis, :])
expected_variance = tf.linalg.diag_part(
tf.reshape(expected_variance, [3, 4, 10, 10]))
expected_variance = tf.linalg.matrix_transpose(
tf.reshape(expected_variance, [3, 4, 5, 2]))
expected_variance = expected_variance * (df /
(df - 2.))[..., tf.newaxis,
tf.newaxis]
mean_, loc_, variance_, expected_variance_ = self.evaluate(
[matrix_t.mean(), loc,
matrix_t.variance(), expected_variance])
self.assertAllClose(loc_, mean_)
self.assertAllClose(expected_variance_, variance_)
