def testTriLAdjoint(self): tril = np.array([[[3, 0, 0], [2, -1, 0], [3, 2, 1]], [[2, 0, 0], [3, -2, 0], [4, 3, 2]]], dtype=np.float32) scale = tf.linalg.LinearOperatorLowerTriangular(tril, is_non_singular=True) bijector = tfb.ScaleMatvecLinearOperator(scale=scale, adjoint=True, validate_args=True) x = np.array([[[1, 0, -1], [2, 3, 4]], [[4, 1, -7], [6, 9, 8]]], dtype=np.float32) # If we made the bijector do x*A+b then this would be simplified to: # y = np.matmul(x, tril). triu = tril.transpose([0, 2, 1]) y = np.matmul(triu, x[..., np.newaxis])[..., 0] ildj = -np.sum(np.log(np.abs(np.diagonal(tril, axis1=-2, axis2=-1)))) self.assertStartsWith(bijector.name, 'scale_matvec_linear_operator') self.assertAllClose(y, self.evaluate(bijector.forward(x))) self.assertAllClose(x, self.evaluate(bijector.inverse(y))) self.assertAllClose( ildj, self.evaluate(bijector.inverse_log_det_jacobian(y, event_ndims=2))) self.assertAllClose( self.evaluate( -bijector.inverse_log_det_jacobian(y, event_ndims=2)), self.evaluate(bijector.forward_log_det_jacobian(x, event_ndims=2)))
def testMean(self): shift = np.array([[-1, 0, 1], [-1, -2, -3]], dtype=np.float32) diag = np.array([[1, 2, 3], [2, 3, 2]], dtype=np.float32) fake_mvn = self._cls()( tfd.MultivariateNormalDiag( loc=tf.zeros_like(shift), scale_diag=tf.ones_like(diag), validate_args=True), tfb.Chain([ tfb.Shift(shift=shift), tfb.ScaleMatvecLinearOperator( scale=tf.linalg.LinearOperatorDiag(diag, is_non_singular=True)) ], validate_args=True), validate_args=True) self.assertAllClose(shift, self.evaluate(fake_mvn.mean()))
def testEntropy(self): shift = np.array([[-1, 0, 1], [-1, -2, -3]], dtype=np.float32) diag = np.array([[1, 2, 3], [2, 3, 2]], dtype=np.float32) actual_mvn_entropy = np.concatenate( [[stats.multivariate_normal(shift[i], np.diag(diag[i]**2)).entropy()] for i in range(len(diag))]) fake_mvn = self._cls()( tfd.MultivariateNormalDiag( loc=tf.zeros_like(shift), scale_diag=tf.ones_like(diag), validate_args=True), tfb.Chain([ tfb.Shift(shift=shift), tfb.ScaleMatvecLinearOperator( scale=tf.linalg.LinearOperatorDiag(diag, is_non_singular=True)) ], validate_args=True), validate_args=True) self.assertAllClose(actual_mvn_entropy, self.evaluate(fake_mvn.entropy()))
def testDiag(self): diag = np.array([[1, 2, 3], [2, 5, 6]], dtype=np.float32) scale = tf.linalg.LinearOperatorDiag(diag, is_non_singular=True) bijector = tfb.ScaleMatvecLinearOperator(scale=scale, validate_args=True) x = np.array([[1, 0, -1], [2, 3, 4]], dtype=np.float32) y = diag * x ildj = -np.sum(np.log(np.abs(diag)), axis=-1) self.assertStartsWith(bijector.name, 'scale_matvec_linear_operator') self.assertAllClose(y, self.evaluate(bijector.forward(x))) self.assertAllClose(x, self.evaluate(bijector.inverse(y))) self.assertAllClose( ildj, self.evaluate(bijector.inverse_log_det_jacobian(y, event_ndims=1))) self.assertAllClose( self.evaluate( -bijector.inverse_log_det_jacobian(y, event_ndims=1)), self.evaluate(bijector.forward_log_det_jacobian(x, event_ndims=1)))