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