def testOneDimNormal(self):
"""Sampling from the Standard Normal Distribution."""
dtype = np.float32
num_results, tolerance = self._get_mode_dependent_settings()
target = tfd.Normal(loc=dtype(0), scale=dtype(1))
samples, _ = tfp.mcmc.sample_chain(
num_results=num_results,
current_state=dtype(1),
kernel=tfp.mcmc.SliceSampler(
target.log_prob,
step_size=1.0,
max_doublings=5,
seed=test_util.test_seed_stream()),
num_burnin_steps=100,
parallel_iterations=1) # For determinism.
sample_mean = tf.reduce_mean(samples, axis=0)
sample_std = tfp.stats.stddev(samples)
self.assertAllClose(0., sample_mean, atol=tolerance, rtol=tolerance)
self.assertAllClose(1., sample_std, atol=tolerance, rtol=tolerance)
def test_specifying_event_shape(
self, event_shape, bijector, dtype, is_static):
seed = test_util.test_seed_stream()
surrogate_posterior = (
tfp.experimental.vi.build_factored_surrogate_posterior(
event_shape=tf.nest.map_structure(
lambda s: self.maybe_static( # pylint: disable=g-long-lambda
np.array(s, dtype=np.int32), is_static=is_static),
event_shape),
bijector=bijector,
initial_unconstrained_loc=functools.partial(
tf.random.uniform, minval=-2., maxval=2., dtype=dtype),
seed=seed(),
validate_args=True))
self.evaluate([v.initializer
for v in surrogate_posterior.trainable_variables])
self._test_shapes(
surrogate_posterior, batch_shape=[], event_shape=event_shape,
seed=seed())
self._test_gradients(surrogate_posterior, seed=seed())
self._test_dtype(surrogate_posterior, dtype, seed())
def testLogBesselKveCorrect(self, dtype, rtol, atol=1e-6):
seed_stream = test_util.test_seed_stream()
v = tf.random.uniform([int(1e5)],
minval=0.1,
maxval=0.5,
seed=seed_stream(),
dtype=dtype)
z = tf.random.uniform([int(1e5)],
minval=1.,
maxval=10.,
seed=seed_stream(),
dtype=dtype)
_, _, log_bessel_kve_expected_, log_bessel_kve_actual_ = self.evaluate(
[
v, z,
tf.math.log(tfp.math.bessel_kve(v, z)),
tfp.math.log_bessel_kve(v, z)
])
self.assertAllClose(log_bessel_kve_expected_,
log_bessel_kve_actual_,
rtol=rtol,
atol=atol)
def make_multivariate_mixture(batch_shape,
num_components,
event_shape,
use_static_graph,
batch_shape_tensor=None):
if batch_shape_tensor is None:
batch_shape_tensor = batch_shape
batch_shape_tensor = tf.convert_to_tensor(value=batch_shape_tensor,
dtype=tf.int32)
seed_stream = test_util.test_seed_stream('multivariate_mixture')
logits = -50. + tf.random.uniform(tf.concat(
(batch_shape_tensor, [num_components]), 0),
-1,
1,
dtype=tf.float32,
seed=seed_stream())
tensorshape_util.set_shape(
logits, tensorshape_util.concatenate(batch_shape, num_components))
static_batch_and_event_shape = (
tf.TensorShape(batch_shape).concatenate(event_shape))
event_shape = tf.convert_to_tensor(value=event_shape, dtype=tf.int32)
batch_and_event_shape = tf.concat((batch_shape_tensor, event_shape), 0)
def create_component():
loc = tf.random.normal(batch_and_event_shape, seed=seed_stream())
scale_diag = 10 * tf.random.uniform(batch_and_event_shape,
seed=seed_stream())
tensorshape_util.set_shape(loc, static_batch_and_event_shape)
tensorshape_util.set_shape(scale_diag, static_batch_and_event_shape)
return tfd.MultivariateNormalDiag(loc=loc,
scale_diag=scale_diag,
validate_args=True)
components = [create_component() for _ in range(num_components)]
cat = tfd.Categorical(logits, dtype=tf.int32)
return tfd.Mixture(cat,
components,
use_static_graph=use_static_graph,
validate_args=True)
def assert_univariate_target_conservation(test, target_d, step_size):
# Sample count limited partly by memory reliably available on Forge. The test
# remains reasonable even if the nuts recursion limit is severely curtailed
# (e.g., 3 or 4 levels), so use that to recover some memory footprint and bump
# the sample count.
num_samples = int(5e4)
num_steps = 1
strm = tfp_test_util.test_seed_stream()
initialization = target_d.sample([num_samples], seed=strm())
@tf.function(autograph=False)
def run_chain():
nuts = tfp.mcmc.NoUTurnSampler(target_d.log_prob,
step_size=step_size,
max_tree_depth=3,
unrolled_leapfrog_steps=2,
seed=strm())
result, _ = tfp.mcmc.sample_chain(num_results=num_steps,
num_burnin_steps=0,
current_state=initialization,
kernel=nuts)
return result
result = run_chain()
test.assertAllEqual([num_steps, num_samples], result.shape)
answer = result[0]
check_cdf_agrees = st.assert_true_cdf_equal_by_dkwm(answer,
target_d.cdf,
false_fail_rate=1e-6)
check_enough_power = assert_util.assert_less(
st.min_discrepancy_of_true_cdfs_detectable_by_dkwm(
num_samples, false_fail_rate=1e-6, false_pass_rate=1e-6), 0.025)
movement = tf.abs(answer - initialization)
test.assertAllEqual([num_samples], movement.shape)
# This movement distance (1 * step_size) was selected by reducing until 100
# runs with independent seeds all passed.
check_movement = assert_util.assert_greater_equal(tf.reduce_mean(movement),
1 * step_size)
return (check_cdf_agrees, check_enough_power, check_movement)
def VerifyPowerSphericalVonMisesFisherZeroKL(self, dim):
seed_stream = test_util.test_seed_stream()
mean_direction = tf.random.uniform(
shape=[5, dim],
minval=self.dtype(1.),
maxval=self.dtype(2.),
dtype=self.dtype,
seed=seed_stream())
mean_direction = tf.nn.l2_normalize(mean_direction, axis=-1)
# Zero concentration is the same as a uniform distribution on the sphere.
# Check that the KL divergence is zero.
concentration = self.dtype(0.)
ps = tfp.distributions.PowerSpherical(
mean_direction=mean_direction,
concentration=concentration)
vmf = tfp.distributions.VonMisesFisher(
mean_direction=mean_direction,
concentration=concentration)
true_kl = tfp.distributions.kl_divergence(ps, vmf)
true_kl_ = self.evaluate(true_kl)
self.assertAllClose(true_kl_, np.zeros_like(true_kl_), atol=1e-4)
def testFourDimNormal(self):
"""Sampling from a 4-D Multivariate Normal distribution."""
dtype = np.float32
true_mean = dtype([0, 4, -8, 2])
true_cov = np.eye(4, dtype=dtype)
num_results, tolerance = _get_mode_dependent_settings()
num_chains = 16
target = tfd.MultivariateNormalTriL(loc=true_mean, scale_tril=true_cov)
# Initial state of the chain
init_state = np.ones([num_chains, 4], dtype=dtype)
# Run Slice Samper for `num_results` iterations for `num_chains`
# independent chains:
states = tfp.mcmc.sample_chain(num_results=num_results,
current_state=init_state,
kernel=tfp.mcmc.SliceSampler(
target_log_prob_fn=tf.function(
target.log_prob,
autograph=False),
step_size=1.0,
max_doublings=5),
num_burnin_steps=100,
trace_fn=None,
seed=test_util.test_seed_stream())
result = tf.reshape(states, [-1, 4])
sample_mean = tf.reduce_mean(result, axis=0)
sample_cov = tfp.stats.covariance(result)
self.assertAllClose(true_mean,
sample_mean,
atol=tolerance,
rtol=tolerance)
self.assertAllClose(true_cov,
sample_cov,
atol=tolerance,
rtol=tolerance)
def test_log_joint_with_missing_observations(self):
# Test that this component accepts MaskedTimeSeries inputs. In most
# cases, it is sufficient that the component accesses only
# `empirical_statistics(observed_time_series)`.
# TODO(b/139483802): De-flake this test when run with --vary_seed.
seed = test_util.test_seed_stream(hardcoded_seed=123)
observed_time_series = np.array(
[1.0, 2.0, -1000., 0.4, np.nan, 1000., 4.2, np.inf]).astype(np.float32)
observation_mask = np.array(
[False, False, True, False, True, True, False, True]).astype(np.bool_)
masked_time_series = tfp.sts.MaskedTimeSeries(observed_time_series,
is_missing=observation_mask)
model = self._build_sts(observed_time_series=masked_time_series)
log_joint_fn = model.joint_log_prob(
observed_time_series=masked_time_series)
lp = self.evaluate(
log_joint_fn(*[p.prior.sample(seed=seed()) for p in model.parameters]))
self.assertEqual(tf.TensorShape([]), lp.shape)
self.assertTrue(np.isfinite(lp))
def test_multipart_bijector(self):
dist = tfd.JointDistributionNamed({
'a': tfd.Exponential(1.),
'b': tfd.Normal(0., 1.),
'c': lambda b, a: tfd.Sample(tfd.Normal(b, a), sample_shape=[5])})
seed = test_util.test_seed_stream()
surrogate_posterior = (
tfp.experimental.vi.build_factored_surrogate_posterior(
event_shape=dist.event_shape,
bijector=(
dist.experimental_default_event_space_bijector()),
initial_unconstrained_loc=functools.partial(
tf.random.uniform, minval=-2., maxval=2.),
seed=seed(),
validate_args=True))
self.evaluate([v.initializer
for v in surrogate_posterior.trainable_variables])
# Test that the posterior has the specified event shape(s).
self.assertAllEqualNested(
self.evaluate(dist.event_shape_tensor()),
self.evaluate(surrogate_posterior.event_shape_tensor()))
posterior_sample_ = self.evaluate(surrogate_posterior.sample(seed=seed()))
posterior_logprob_ = self.evaluate(
surrogate_posterior.log_prob(posterior_sample_))
# Test that all sample Tensors have the expected shapes.
check_shape = lambda s, x: self.assertAllEqual(s, x.shape)
self.assertAllAssertsNested(
check_shape, dist.event_shape, posterior_sample_)
# Test that samples are finite and not NaN.
self.assertAllAssertsNested(self.assertAllFinite, posterior_sample_)
# Test that logprob is scalar, finite, and not NaN.
self.assertEmpty(posterior_logprob_.shape)
self.assertAllFinite(posterior_logprob_)
def test_kahan_precision(self, jit=False):
maybe_jit = lambda f: f
if jit:
self.skip_if_no_xla()
maybe_jit = tf.function(jit_compile=True)
n = 20_000
stream = test_util.test_seed_stream()
samps = tfd.Normal(0, 1).sample(n, seed=stream())
scale = tfd.LogNormal(0, .2).sample([7, 1], seed=stream())
mvn = tfd.MultivariateNormalDiag(
loc=tf.zeros([n]), scale_diag=tf.zeros([n]) + scale,
experimental_use_kahan_sum=True)
mvn64 = tfd.MultivariateNormalDiag(
loc=tf.zeros([n], dtype=tf.float64),
scale_diag=tf.zeros([n], dtype=tf.float64) + tf.cast(scale, tf.float64))
lp = maybe_jit(mvn.log_prob)(samps)
lp64 = mvn64.log_prob(tf.cast(samps, tf.float64))
lp, lp64 = self.evaluate((tf.cast(lp, tf.float64), lp64))
# Without fastmath, without Kahan fails ~30-100%, max abs error up to .2
self.assertAllClose(lp64, lp, rtol=0, atol=.006)
def testOwensTOddEven(self, dtype):
seed_stream = test_util.test_seed_stream()
a = tf.random.uniform(shape=[int(1e3)],
minval=0.,
maxval=100.,
dtype=dtype,
seed=seed_stream())
h = tf.random.uniform(shape=[int(1e3)],
minval=0.,
maxval=100.,
dtype=dtype,
seed=seed_stream())
# OwensT(h, a) = OwensT(-h, a)
self.assertAllClose(
self.evaluate(tfp.math.owens_t(h, a)),
self.evaluate(tfp.math.owens_t(-h, a)),
)
# OwensT(h, a) = -OwensT(h, -a)
self.assertAllClose(
self.evaluate(tfp_math.owens_t(h, a)),
self.evaluate(-tfp_math.owens_t(h, -a)),
)
def test_explicit_init_samples(self):
stream = test_util.test_seed_stream()
# Compute everything in a function so it is consistent in graph mode
@tf.function
def do_sample():
jd_model = tfd.JointDistributionNamed({
'x': tfd.HalfNormal(1.),
'y': lambda x: tfd.Normal(0., x)})
init = {'x': tf.ones(64)}
return tfp.experimental.mcmc.windowed_adaptive_hmc(
10,
jd_model,
num_adaptation_steps=200,
current_state=init,
num_leapfrog_steps=5,
discard_tuning=False,
y=tf.constant(1.),
seed=stream(),
trace_fn=None)
self.evaluate(do_sample())
def test_kahan_precision(self, jit=False):
maybe_jit = lambda f: f
if jit:
self.skip_if_no_xla()
maybe_jit = tf.function(experimental_compile=True)
stream = test_util.test_seed_stream()
n = 20_000
samps = tfd.Poisson(rate=1.).sample(n, seed=stream())
log_rate = tf.fill([n], tfd.Normal(0, .2).sample(seed=stream()))
pois = tfd.Poisson(log_rate=log_rate)
lp_fn = maybe_jit(
tfd.Independent(pois,
reinterpreted_batch_ndims=1,
experimental_use_kahan_sum=True).log_prob)
lp = lp_fn(samps)
pois64 = tfd.Poisson(log_rate=tf.cast(log_rate, tf.float64))
lp64 = tfd.Independent(pois64, reinterpreted_batch_ndims=1).log_prob(
tf.cast(samps, tf.float64))
# Evaluate together to ensure we use the same samples.
lp, lp64 = self.evaluate((tf.cast(lp, tf.float64), lp64))
# Fails ~75% CPU, 1-75% GPU --vary_seed runs w/o experimental_use_kahan_sum.
self.assertAllClose(lp64, lp, rtol=0., atol=.01)
def testMVNConjugateLinearUpdatePreservesStructuredLinops(self):
strm = test_util.test_seed_stream()
num_outputs = 4
prior_scale = tf.linalg.LinearOperatorScaledIdentity(num_outputs, 4.)
likelihood_scale = tf.linalg.LinearOperatorScaledIdentity(num_outputs, 0.2)
linear_transformation = tf.linalg.LinearOperatorIdentity(num_outputs)
observation = tf.random.normal([num_outputs], seed=strm())
posterior_mean, posterior_prec = (
tfd.mvn_conjugate_linear_update(
prior_scale=prior_scale,
linear_transformation=linear_transformation,
likelihood_scale=likelihood_scale,
observation=observation))
# TODO(davmre): enable next line once internal CI is updated to recent TF.
# self.assertIsInstance(posterior_prec,
# tf.linalg.LinearOperatorScaledIdentity)
self._mvn_linear_update_test_helper(
prior_mean=tf.zeros([num_outputs]),
prior_scale=prior_scale.to_dense(),
linear_transformation=linear_transformation.to_dense(),
likelihood_scale=likelihood_scale.to_dense(),
observation=observation,
candidate_posterior_mean=posterior_mean,
candidate_posterior_prec=posterior_prec.to_dense())
# Also check the result against the scalar calculation.
scalar_posterior_dist = tfd.normal_conjugates_known_scale_posterior(
prior=tfd.Normal(loc=0., scale=prior_scale.diag_part()),
scale=likelihood_scale.diag_part(),
s=observation, n=1)
(posterior_mean_, posterior_prec_,
scalar_posterior_mean_, scalar_posterior_prec_) = self.evaluate(
(posterior_mean, posterior_prec.to_dense(),
scalar_posterior_dist.mean(),
tf.linalg.diag(1./scalar_posterior_dist.variance())))
self.assertAllClose(posterior_mean_, scalar_posterior_mean_)
self.assertAllClose(posterior_prec_, scalar_posterior_prec_)
def test_poisson_switchover_graphical_model(self):
# Build a pretend dataset.
seed = test_util.test_seed_stream(salt='poisson')
n = [43, 31]
count_data = tf.cast(tf.concat([
tfd.Poisson(rate=15.).sample(n[0], seed=seed()),
tfd.Poisson(rate=25.).sample(n[1], seed=seed()),
],
axis=0),
dtype=tf.float32)
count_data = self.evaluate(count_data)
n = np.sum(n)
# Make model.
gather = lambda tau, lambda_: tf.gather( # pylint: disable=g-long-lambda
lambda_,
indices=tf.cast(tau[..., tf.newaxis] < tf.linspace(0., 1., n),
dtype=tf.int32),
# TODO(b/139204153): Remove static value hack after bug closed.
batch_dims=int(tf.get_static_value(tf.rank(tau))))
alpha = tf.math.reciprocal(tf.reduce_mean(count_data))
joint = tfd.JointDistributionSequential(
[
tfd.Sample(tfd.Exponential(rate=alpha), sample_shape=[2]),
tfd.Uniform(),
lambda tau, lambda_: tfd.Independent( # pylint: disable=g-long-lambda
tfd.Poisson(rate=gather(tau, lambda_)),
reinterpreted_batch_ndims=1),
],
validate_args=True)
# Verify model correctly "compiles".
batch_shape = [3, 4]
self.assertEqual(
batch_shape,
joint.log_prob(
joint.sample(batch_shape, seed=test_util.test_seed())).shape)
def testSampleAgainstProb(self):
seed_stream = test_util.test_seed_stream()
n = 4
c1 = self.evaluate(1. + 2. * tf.random.uniform(
shape=[4, 3, 2], dtype=tf.float32, seed=seed_stream()))
c0 = self.evaluate(1. + 2. * tf.random.uniform(
shape=[4, 3, 1], dtype=tf.float32, seed=seed_stream()))
dist = tfd.BetaBinomial(tf.cast(n, dtype=tf.float32),
c1,
c0,
validate_args=True)
num_samples = int(1e4)
x = self.evaluate(
tf.cast(dist.sample(num_samples, seed=seed_stream()), tf.int32))
for i in range(n + 1):
self.assertAllClose(self.evaluate(dist.prob(i)),
np.sum(x == i, axis=0) / (num_samples * 1.0),
atol=0.01,
rtol=0.1)
def test_specifying_initial_loc(self, event_shape, initial_loc,
implicit_batch_shape, bijector, dtype,
is_static):
initial_loc = tf.nest.map_structure(
lambda s: self.maybe_static( # pylint: disable=g-long-lambda
np.array(s, dtype=dtype),
is_static=is_static),
initial_loc)
if bijector is not None:
initial_unconstrained_loc = tf.nest.map_structure(
lambda x, b: x
if b is None else b.inverse(x), initial_loc, bijector)
else:
initial_unconstrained_loc = initial_loc
surrogate_posterior = (
tfp.experimental.vi.build_factored_surrogate_posterior(
event_shape=event_shape,
initial_unconstrained_loc=initial_unconstrained_loc,
initial_unconstrained_scale=1e-6,
bijector=bijector,
validate_args=True))
self.evaluate(
[v.initializer for v in surrogate_posterior.trainable_variables])
seed = test_util.test_seed_stream()
self._test_shapes(surrogate_posterior,
batch_shape=implicit_batch_shape,
event_shape=event_shape,
seed=seed())
self._test_gradients(surrogate_posterior, seed=seed())
self._test_dtype(surrogate_posterior, dtype, seed())
# Check that the sampled values are close to the initial locs.
posterior_sample_ = self.evaluate(
surrogate_posterior.sample(seed=seed()))
self.assertAllCloseNested(initial_loc, posterior_sample_, atol=1e-4)
def test_specifying_event_shape(
self, event_shape, bijector, dtype, is_static):
seed = test_util.test_seed_stream()
surrogate_posterior = (
tfp.experimental.vi.build_factored_surrogate_posterior(
event_shape=tf.nest.map_structure(
lambda s: self.maybe_static( # pylint: disable=g-long-lambda
np.array(s, dtype=np.int32), is_static=is_static),
event_shape),
bijector=bijector,
initial_unconstrained_loc=functools.partial(
tf.random.uniform, minval=-2., maxval=2., dtype=dtype),
seed=seed(),
validate_args=True))
self.evaluate([v.initializer
for v in surrogate_posterior.trainable_variables])
posterior_sample_ = self.evaluate(surrogate_posterior.sample(seed=seed()))
posterior_logprob_ = self.evaluate(
surrogate_posterior.log_prob(posterior_sample_))
posterior_event_shape = self.evaluate(
surrogate_posterior.event_shape_tensor())
# Test that the posterior has the specified event shape(s).
self.assertAllEqualNested(event_shape, posterior_event_shape)
# Test that all sample Tensors have the expected shapes.
check_shape = lambda s, x: self.assertAllEqual(s, x.shape)
tf.nest.map_structure(check_shape, event_shape, posterior_sample_)
self.assertAllEqual([], posterior_logprob_.shape)
# Test that gradients are available wrt the variational parameters.
self.assertNotEmpty(surrogate_posterior.trainable_variables)
with tf.GradientTape() as tape:
posterior_logprob = surrogate_posterior.log_prob(posterior_sample_)
grad = tape.gradient(posterior_logprob,
surrogate_posterior.trainable_variables)
self.assertTrue(all(g is not None for g in grad))
def testParamPropertiesAndFromParams(self):
classes = [
tfd.Normal,
tfd.Bernoulli,
tfd.Beta,
tfd.Chi2,
tfd.Exponential,
tfd.Gamma,
tfd.InverseGamma,
tfd.Laplace,
tfd.StudentT,
tfd.Uniform,
]
sample_shapes = [(), (10, ), (10, 20, 30)]
seed_stream = test_util.test_seed_stream('param_shapes')
for cls in classes:
for sample_shape in sample_shapes:
parameter_properties = cls.parameter_properties()
param_shapes = {
name: param.shape_fn(sample_shape) # pylint: disable=comprehension-too-complex
for name, param in parameter_properties.items()
if param.is_preferred
}
params = {
name: tf.random.normal(shape, seed=seed_stream())
for name, shape in param_shapes.items()
}
dist = cls(**params)
self.assertAllEqual(
sample_shape,
self.evaluate(tf.shape(dist.sample(seed=seed_stream()))))
dist_copy = dist.copy()
self.assertAllEqual(
sample_shape,
self.evaluate(
tf.shape(dist_copy.sample(seed=seed_stream()))))
self.assertEqual(dist.parameters, dist_copy.parameters)
def testLogProbAgainstDirichletMultinomial(self):
seed_stream = test_util.test_seed_stream()
n = tf.constant([10., 20., 30.])
c1 = self.evaluate(1. + 2. * tf.random.uniform(
shape=[4, 3], dtype=tf.float32, seed=seed_stream()))
c0 = self.evaluate(1. + 2. * tf.random.uniform(
shape=[4, 3], dtype=tf.float32, seed=seed_stream()))
beta_binomial = tfd.BetaBinomial(n, c1, c0, validate_args=True)
dirichlet_multinomial = tfd.DirichletMultinomial(n,
tf.stack([c1, c0],
axis=-1),
validate_args=True)
num_samples = 3
beta_binomial_sample = self.evaluate(
beta_binomial.sample(num_samples, seed=seed_stream()))
beta_binomial_log_prob = beta_binomial.log_prob(beta_binomial_sample)
dirichlet_multinomial_log_prob = dirichlet_multinomial.log_prob(
tf.stack([beta_binomial_sample, n - beta_binomial_sample],
axis=-1))
self.assertAllClose(self.evaluate(beta_binomial_log_prob),
self.evaluate(dirichlet_multinomial_log_prob),
rtol=1e-4,
atol=1e-4)
dirichlet_multinomial_sample = self.evaluate(
dirichlet_multinomial.sample(num_samples, seed=seed_stream()))
dirichlet_multinomial_log_prob = dirichlet_multinomial.log_prob(
dirichlet_multinomial_sample)
beta_binomial_log_prob = beta_binomial.log_prob(
tf.squeeze(dirichlet_multinomial_sample[..., 0]))
self.assertAllClose(self.evaluate(dirichlet_multinomial_log_prob),
self.evaluate(beta_binomial_log_prob),
rtol=1e-4,
atol=1e-4)
def test_random_projections(self, dtype):
strm = tfp_test_util.test_seed_stream()
rng = np.random.RandomState(seed=strm() % 2**31)
num_samples = 57000
# Validate experiment design
# False fail rate here is the target rate of 1e-6 divided by the number of
# projections.
d = st.min_discrepancy_of_true_cdfs_detectable_by_dkwm_two_sample(
num_samples, num_samples, false_fail_rate=1e-8, false_pass_rate=1e-6)
# Choose num_samples so the discrepancy is below 0.05, which should be
# enough to detect a mean shift of around 1/8 of a standard deviation, or a
# scale increase of around 25% (in any particular projection).
self.assertLess(self.evaluate(d), 0.05)
ground_truth = rng.multivariate_normal(
mean=[0, 0], cov=[[1, 0.5], [0.5, 1]], size=num_samples).astype(dtype)
more_samples = rng.multivariate_normal(
mean=[0, 0], cov=[[1, 0.5], [0.5, 1]], size=num_samples).astype(dtype)
self.evaluate(
st.assert_multivariate_true_cdf_equal_on_projections_two_sample(
ground_truth, more_samples, num_projections=100,
false_fail_rate=1e-6, seed=strm()))
def assert_catches_mistake(mean, cov):
wrong_samples = rng.multivariate_normal(
mean=mean, cov=cov, size=num_samples).astype(dtype=dtype)
msg = 'Empirical CDFs outside joint K-S envelope'
with self.assertRaisesOpError(msg):
self.evaluate(
st.assert_multivariate_true_cdf_equal_on_projections_two_sample(
ground_truth, wrong_samples, num_projections=100,
false_fail_rate=1e-6, seed=strm()))
assert_catches_mistake([0, 1], [[1, 0.5], [0.5, 1]])
assert_catches_mistake([0, 0], [[1, 0.7], [0.7, 1]])
assert_catches_mistake([0, 0], [[1, 0.3], [0.3, 1]])
def test_ensemble_kalman_filter_constant_model_multivariate(self):
def transition_fn(_, particles, extra):
return tfd.MultivariateNormalDiag(loc=particles,
scale_diag=[1e-11] * 2), extra
def observation_fn(_, particles, extra):
return tfd.MultivariateNormalDiag(loc=particles,
scale_diag=[1e-1] * 2), extra
seed_stream = test_util.test_seed_stream()
# Initialize the ensemble.
particles = self.evaluate(
tf.random.normal(shape=[300, 3, 2],
seed=seed_stream(),
dtype=tf.float64))
state = tfs.EnsembleKalmanFilterState(step=0,
particles=particles,
extra={'unchanged': 1})
for _ in range(8):
state = tfs.ensemble_kalman_filter_predict(
state,
transition_fn=transition_fn,
seed=seed_stream(),
inflate_fn=None)
state = tfs.ensemble_kalman_filter_update(
state,
observation=[0., 0.],
observation_fn=observation_fn,
seed=seed_stream())
self.assertAllClose([[0., 0.]] * 3,
self.evaluate(
tf.reduce_mean(state.particles, axis=0)),
atol=1e-2)
def VerifyPowerSphericalVonMisesFisherKL(self, dim):
seed_stream = test_util.test_seed_stream()
mean_direction1 = tf.random.uniform(shape=[5, dim],
minval=self.dtype(1.),
maxval=self.dtype(2.),
dtype=self.dtype,
seed=seed_stream())
mean_direction2 = tf.random.uniform(shape=[5, dim],
minval=self.dtype(1.),
maxval=self.dtype(2.),
dtype=self.dtype,
seed=seed_stream())
mean_direction1 = tf.nn.l2_normalize(mean_direction1, axis=-1)
mean_direction2 = tf.nn.l2_normalize(mean_direction2, axis=-1)
concentration1 = tf.math.log(
tf.random.uniform(shape=[2, 1],
minval=self.dtype(1.),
maxval=self.dtype(100.),
dtype=self.dtype,
seed=seed_stream()))
concentration2 = tf.math.log(
tf.random.uniform(shape=[2, 1],
minval=self.dtype(1.),
maxval=self.dtype(100.),
dtype=self.dtype,
seed=seed_stream()))
ps = tfp.distributions.PowerSpherical(mean_direction=mean_direction1,
concentration=concentration1)
vmf = tfp.distributions.VonMisesFisher(mean_direction=mean_direction2,
concentration=concentration2)
x = ps.sample(int(6e4), seed=test_util.test_seed())
kl_sample = tf.reduce_mean(ps.log_prob(x) - vmf.log_prob(x), axis=0)
true_kl = tfp.distributions.kl_divergence(ps, vmf)
true_kl_, kl_sample_ = self.evaluate([true_kl, kl_sample])
self.assertAllClose(true_kl_, kl_sample_, atol=0.0, rtol=7e-2)
def VerifyVonMisesFisherUniformKL(self, dim):
seed_stream = test_util.test_seed_stream()
mean_direction = tf.random.uniform(shape=[4, dim],
minval=1.,
maxval=4.,
seed=seed_stream())
mean_direction = tf.nn.l2_normalize(mean_direction, axis=-1)
concentration = tf.math.log(
tf.random.uniform(shape=[2, 1],
minval=2.,
maxval=20.,
seed=seed_stream()))
vmf = tfp.distributions.VonMisesFisher(mean_direction=mean_direction,
concentration=concentration)
su = tfp.distributions.SphericalUniform(dimension=dim)
x = vmf.sample(int(5e4), seed=test_util.test_seed())
kl_sample = tf.reduce_mean(vmf.log_prob(x) - su.log_prob(x), axis=0)
true_kl = tfp.distributions.kl_divergence(vmf, su)
true_kl_, kl_sample_ = self.evaluate([true_kl, kl_sample])
self.assertAllClose(true_kl_, kl_sample_, atol=0.0, rtol=0.3)
def test_step_size_changed(self):
target_dist = tfd.MultivariateNormalDiag(loc=[0., 0.],
scale_diag=[1., 10.])
# `hmc_kernel`'s step size is far from optimal
hmc_kernel = tfp.mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=target_dist.log_prob,
num_leapfrog_steps=27,
step_size=10)
step_adaptation_kernel = tfp.mcmc.SimpleStepSizeAdaptation(
inner_kernel=hmc_kernel,
adaptation_rate=0.8,
num_adaptation_steps=9)
trans_kernel = tfp.mcmc.TransformedTransitionKernel(
inner_kernel=step_adaptation_kernel, bijector=tfb.Exp())
kernel_results = trans_kernel.inner_kernel.bootstrap_results(
tf.zeros(2))
stream = test_util.test_seed_stream()
for _ in range(2):
_, kernel_results = trans_kernel.inner_kernel.one_step(
tf.zeros(2), kernel_results, seed=stream())
adapted_step_size = self.evaluate(
kernel_results.inner_results.accepted_results.step_size)
self.assertLess(adapted_step_size, 7)
def testLargeLogProbDiffScalarUnderlying(self):
shp = [25, 200]
d0 = tfd.Sample(tfd.Normal(0., .1), shp)
d1 = tfd.Sample(tfd.Normal(1e-5, .1), shp)
strm = test_util.test_seed_stream()
x0 = self.evaluate( # overdispersed
tfd.Normal(0, 2).sample(shp, seed=strm()))
x1 = self.evaluate( # overdispersed, perturbed
x0 + tfd.Normal(0, 1e-6).sample(x0.shape, seed=strm()))
d0_64 = d0.copy(distribution=tfd.Normal(
tf.cast(d0.distribution.loc, tf.float64),
tf.cast(d0.distribution.scale, tf.float64)))
d1_64 = d1.copy(distribution=tfd.Normal(
tf.cast(d1.distribution.loc, tf.float64),
tf.cast(d1.distribution.scale, tf.float64)))
oracle_64 = tf.reduce_sum(
d0_64.distribution.log_prob(tf.cast(x0, tf.float64)) -
d1_64.distribution.log_prob(tf.cast(x1, tf.float64)))
self.assertAllClose(oracle_64,
tfp.experimental.distributions.log_prob_ratio(
d0, x0, d1, x1),
rtol=0.,
atol=0.007)
def test_recip_scale_exp(self):
p = tfb.Reciprocal()(tfb.Scale(3.)(tfb.Exp()))
stream = test_util.test_seed_stream()
dim = 2
x = self.evaluate(tf.random.uniform([4, dim], seed=stream()))
q = tfb.Reciprocal()(tfb.Scale(2.)(tfb.Exp()))
y = self.evaluate(tf.random.uniform([4, dim], seed=stream()))
expected_fldjr = ((x - y) + # Exp.fldj
(np.log(3) - np.log(2)) + # Scale.fldj
-2 * (np.log(3. * np.exp(x) /
(2 * np.exp(y)))) # Reciprocal.fldj
).sum(-1)
self.assertAllClose(
expected_fldjr,
tfeb.forward_log_det_jacobian_ratio(p, x, q, y, event_ndims=1))
self.assertAllClose(
expected_fldjr,
p.forward_log_det_jacobian(x, 1) -
q.forward_log_det_jacobian(y, 1))
self.assertAllClose(
p.inverse_log_det_jacobian(x, 1) -
q.inverse_log_det_jacobian(y, 1),
tfeb.inverse_log_det_jacobian_ratio(p, x, q, y, event_ndims=1))
def testPairwiseSquareDistanceMatrix(self, feature_ndims, dims):
batch_shape = [2, 3]
seed_stream = test_util.test_seed_stream('pairwise_square_distance')
x1 = tf.random.normal(dtype=np.float64,
shape=batch_shape + [dims] * feature_ndims,
seed=seed_stream())
x2 = tf.random.normal(dtype=np.float64,
shape=batch_shape + [dims] * feature_ndims,
seed=seed_stream())
pairwise_square_distance = util.pairwise_square_distance_matrix(
x1, x2, feature_ndims)
x1_pad = util.pad_shape_with_ones(x1,
ndims=1,
start=-(feature_ndims + 1))
x2_pad = util.pad_shape_with_ones(x2,
ndims=1,
start=-(feature_ndims + 2))
actual_square_distance = util.sum_rightmost_ndims_preserving_shape(
tf.math.squared_difference(x1_pad, x2_pad), feature_ndims)
pairwise_square_distance_, actual_square_distance_ = self.evaluate(
[pairwise_square_distance, actual_square_distance])
self.assertAllClose(pairwise_square_distance_, actual_square_distance_)
def VerifyEntropy(self, dim):
seed_stream = test_util.test_seed_stream()
mean_direction = tf.random.uniform(shape=[5, dim],
minval=self.dtype(1.),
maxval=self.dtype(2.),
dtype=self.dtype,
seed=seed_stream())
mean_direction = tf.nn.l2_normalize(mean_direction, axis=-1)
concentration = tf.math.log(
tf.random.uniform(shape=[2, 1],
minval=self.dtype(1.),
maxval=self.dtype(100.),
dtype=self.dtype,
seed=seed_stream()))
ps = tfp.distributions.PowerSpherical(mean_direction=mean_direction,
concentration=concentration,
validate_args=True,
allow_nan_stats=False)
samples = ps.sample(int(3e4), seed=test_util.test_seed())
sample_entropy = -tf.reduce_mean(ps.log_prob(samples), axis=0)
true_entropy, sample_entropy = self.evaluate(
[ps.entropy(), sample_entropy])
self.assertAllClose(sample_entropy, true_entropy, rtol=3e-2)
def test_default_priors_follow_batch_shapes(self):
seed = test_util.test_seed_stream()
num_timesteps = 3
time_series_sample_shape = [4, 2]
observation_shape_full = time_series_sample_shape + [num_timesteps]
dummy_observation = np.random.randn(
*(observation_shape_full)).astype(np.float32)
model = self._build_sts(observed_time_series=dummy_observation)
# The model should construct a default parameter prior for *each* observed
# time series, so the priors will have batch_shape equal to
# `time_series_sample_shape`.
for parameter in model.parameters:
self.assertEqual(parameter.prior.batch_shape, time_series_sample_shape)
# The initial state prior should also have the appropriate batch shape.
# To test this, we build the ssm and test that it has a consistent
# broadcast batch shape.
param_samples = [p.prior.sample(seed=seed()) for p in model.parameters]
ssm = model.make_state_space_model(
num_timesteps=num_timesteps, param_vals=param_samples)
self.assertEqual(ssm.batch_shape, time_series_sample_shape)
