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