def testMVNConjugateLinearUpdateSupportsBatchShape(self): strm = test_util.test_seed_stream() num_latents = 2 num_outputs = 4 batch_shape = [3, 1] prior_mean = tf.ones([num_latents]) prior_scale = tf.eye(num_latents) * 5. likelihood_scale = tf.linalg.LinearOperatorLowerTriangular( tfb.ScaleTriL().forward( tf.random.normal(shape=batch_shape + [int(num_outputs * (num_outputs + 1) / 2)], seed=strm()))) linear_transformation = tf.random.normal( batch_shape + [num_outputs, num_latents], seed=strm()) * 5. true_latent = tf.random.normal(batch_shape + [num_latents], seed=strm()) observation = tf.linalg.matvec(linear_transformation, true_latent) posterior_mean, posterior_prec = (tfd.mvn_conjugate_linear_update( prior_mean=prior_mean, prior_scale=prior_scale, linear_transformation=linear_transformation, likelihood_scale=likelihood_scale, observation=observation)) self._mvn_linear_update_test_helper( prior_mean=prior_mean, prior_scale=prior_scale, linear_transformation=linear_transformation, likelihood_scale=likelihood_scale.to_dense(), observation=observation, candidate_posterior_mean=posterior_mean, candidate_posterior_prec=posterior_prec.to_dense())
def new(params, event_size, validate_args=False, name=None): """Create the distribution instance from a `params` vector.""" with tf.name_scope(name, 'MultivariateNormalTriL', [params, event_size]): return tfd.MultivariateNormalTriL( loc=params[..., :event_size], scale_tril=tfb.ScaleTriL(validate_args=validate_args)( params[..., event_size:]), validate_args=validate_args)
def new(params, event_size, validate_args=False, name=None): """Create the distribution instance from a `params` vector.""" with tf.name_scope(name, 'MultivariateNormalTriL', [params, event_size]): params = tf.convert_to_tensor(params, name='params') scale_tril = tfb.ScaleTriL(diag_shift=np.array( 1e-5, params.dtype.as_numpy_dtype()), validate_args=validate_args) return tfd.MultivariateNormalTriL(loc=params[..., :event_size], scale_tril=scale_tril( params[..., event_size:]), validate_args=validate_args)
def testComputesCorrectValues(self): shift = 1.61803398875 x = np.float32(np.array([-1, .5, 2])) y = np.float32( np.array([[np.exp(2) + shift, 0.], [.5, np.exp(-1) + shift]])) b = tfb.ScaleTriL(diag_bijector=tfb.Exp(), diag_shift=shift) y_ = self.evaluate(b.forward(x)) self.assertAllClose(y, y_, rtol=1e-4) x_ = self.evaluate(b.inverse(y)) self.assertAllClose(x, x_, rtol=1e-4)
def testInvertible(self): # Generate random inputs from an unconstrained space, with # event size 6 to specify 3x3 triangular matrices. batch_shape = [2, 1] x = np.float32(self._rng.randn(*(batch_shape + [6]))) b = tfb.ScaleTriL(diag_bijector=tfb.Softplus(), diag_shift=3.14159) y = self.evaluate(b.forward(x)) self.assertAllEqual(y.shape, batch_shape + [3, 3]) x_ = self.evaluate(b.inverse(y)) self.assertAllClose(x, x_, rtol=1e-4) fldj = self.evaluate(b.forward_log_det_jacobian(x, event_ndims=1)) ildj = self.evaluate(b.inverse_log_det_jacobian(y, event_ndims=2)) self.assertAllClose(fldj, -ildj, rtol=1e-4)
def testJacobian(self): cholesky_to_vector = tfb.Invert( tfb.ScaleTriL(diag_bijector=tfb.Exp(), diag_shift=None)) bijector = tfb.CholeskyToInvCholesky() for x in [ np.array([[2.]], dtype=np.float64), np.array([[2., 0.], [3., 4.]], dtype=np.float64), np.array([[2., 0., 0.], [3., 4., 0.], [5., 6., 7.]], dtype=np.float64) ]: fldj = bijector.forward_log_det_jacobian(x, event_ndims=2) fldj_numerical = self._get_fldj_numerical( bijector, x, event_ndims=2, input_to_vector=cholesky_to_vector, output_to_vector=cholesky_to_vector) fldj_, fldj_numerical_ = self.evaluate([fldj, fldj_numerical]) self.assertAllClose(fldj_, fldj_numerical_, rtol=1e-2)
def new(params, event_size, covariance_type, softplus_scale, validate_args=False, name=None): """Create the distribution instance from a `params` vector.""" covariance_type = str(covariance_type).lower().strip() assert covariance_type in ('full', 'tril', 'diag'), \ "No support for given covariance_type: '%s'" % covariance_type scale_fn = lambda x: tf.math.softplus(x) + tfd.softplus_inverse(1.0) \ if bool(softplus_scale) else x with tf.compat.v1.name_scope(name, 'MultivariateNormal', [params, event_size]): params = tf.convert_to_tensor(value=params, name='params') if covariance_type == 'tril': scale_tril = tfb.ScaleTriL(diag_shift=np.array( 1e-5, params.dtype.as_numpy_dtype()), validate_args=validate_args) return tfd.MultivariateNormalTriL( loc=params[..., :event_size], scale_tril=scale_tril(scale_fn(params[..., event_size:])), validate_args=validate_args) elif covariance_type == 'diag': return tfd.MultivariateNormalDiag( loc=params[..., :event_size], scale_diag=scale_fn(params[..., event_size:])) elif covariance_type == 'full': return tfd.MultivariateNormalFullCovariance( loc=params[..., :event_size], covariance_matrix=tf.reshape( scale_fn(params[..., event_size:]), (event_size, event_size)))