def test_tf_while_var_with_dynamic_shape(self):
rng = test_util.test_np_rng()
x = rng.rand(100, 10)
tensor_x = tf.convert_to_tensor(x, dtype=tf.float32)
running_var = tfp.experimental.stats.RunningVariance((10, ))
def _loop_body(i, state):
if not tf.executing_eagerly():
sample = tf1.placeholder_with_default(tensor_x[i], shape=None)
else:
sample = tensor_x[i]
return (i + 1, running_var.update(state, sample))
_, state = tf.while_loop(
lambda i, _: i < 100,
_loop_body,
(tf.constant(0, dtype=tf.int32), running_var.initialize()),
shape_invariants=(None,
tfp.experimental.stats.RunningCovarianceState(
None, tf.TensorShape(None),
tf.TensorShape(None))))
final_mean, final_var = self.evaluate(
[state.mean, running_var.finalize(state)])
self.assertEqual(final_mean.shape, (10, ))
self.assertEqual(final_var.shape, (10, ))
self.assertAllClose(final_var, np.var(x, axis=0), rtol=1e-5)
def test_windowed_mean_corner_cases(self):
rng = test_util.test_np_rng()
x = rng.rand(7)
# Test mean of an empty set
mean = tfp.stats.windowed_mean(x,
low_indices=[4],
high_indices=4,
axis=0)
mean = self.evaluate(mean)
self.assertAllClose(mean, tf.zeros_like(mean))
# Test mean of a "negative" set. It's the same
# as the mean of the same set spelled "positively", but we
# need to be careful about the inclusive/exclusive semantics of
# the indices.
mean_neg = tfp.stats.windowed_mean(x,
low_indices=[3, 5],
high_indices=[1, 2])
mean_pos = tfp.stats.windowed_mean(x,
low_indices=[1, 2],
high_indices=[3, 5])
mean_neg, mean_pos = self.evaluate([mean_neg, mean_pos])
self.assertAllClose(mean_neg, mean_pos)
# Test default windows: [0, 1), [1, 2), [1, 3), [2, 4), etc
y = [0., 1., 2., 3.]
self.assertAllClose([0., 1., 1.5, 2.5],
self.evaluate(tfp.stats.windowed_mean(y)))
def test_step_function_sequence(self):
if tf.executing_eagerly() and not self.use_static_shape:
# TODO(b/122840816): Modify this test so that it runs in eager mode with
# dynamic shapes, or document that this is the intended behavior.
return
rng = test_util.test_np_rng()
# x jumps to new random value every 10 steps. So correlation length = 10.
x = (rng.randint(-10, 10, size=(1000, 1)) * np.ones(
(1, 10))).ravel().astype(self.dtype)
x_ph = tf1.placeholder_with_default(
x, shape=(1000 * 10, ) if self.use_static_shape else None)
rxx = tfp.stats.auto_correlation(x_ph,
max_lags=1000 * 10 // 2,
center=True,
normalize=False)
if self.use_static_shape:
self.assertAllEqual((1000 * 10 // 2 + 1, ), rxx.shape)
rxx_ = self.evaluate(rxx)
rxx_ /= rxx_[0]
# Expect positive correlation for the first 10 lags, then significantly
# smaller negative.
self.assertGreater(rxx_[:10].min(), 0)
# TODO(b/138375951): Re-enable this assertion once we know why its
# failing.
# self.assertGreater(rxx_[9], 5 * rxx_[10:20].mean())
# RXX should be decreasing for the first 10 lags.
diff = np.diff(rxx_)
self.assertLess(diff[:10].max(), 0)
def test_dynamic_shape_running_covariance(self):
rng = test_util.test_np_rng()
x = rng.rand(100, 3, 5, 2)
running_cov = tfp.experimental.stats.RunningCovariance(tf.TensorShape(
[5, 2]),
event_ndims=1)
state = running_cov.initialize()
for sample in x:
if not tf.executing_eagerly():
sample = tf1.placeholder_with_default(sample, shape=None)
state = running_cov.update(state, sample, axis=0)
final_mean, final_cov = self.evaluate(
[state.mean, running_cov.finalize(state)])
self.assertAllClose(np.mean(x.reshape(300, 5, 2), axis=0),
final_mean,
rtol=1e-5)
self.assertEqual(final_cov.shape, (5, 2, 2))
# manual computation with loops
manual_cov = np.zeros((5, 2, 2))
x_reshaped = x.reshape((300, 5, 2))
delta_mean = x_reshaped - np.mean(x_reshaped, axis=0)
for residual in delta_mean:
for i in range(5):
for j in range(2):
for k in range(2):
manual_cov[i][j][k] += residual[i][j] * residual[i][k]
manual_cov /= 300
self.assertAllClose(manual_cov, final_cov, rtol=1e-5)
def test_batch_vector_sampaxis02_eventaxis1(self):
rng = test_util.test_np_rng()
# x.shape = sample, event, sample, batch
x = rng.randn(2, 3, 4, 5)
y = x + 0.1 * rng.randn(2, 3, 4, 5)
cov = tfp.stats.covariance(x, y, sample_axis=[0, 2], event_axis=[1])
self.assertAllEqual((3, 3, 5), cov.shape)
cov = self.evaluate(cov)
cov_kd = tfp.stats.covariance(x,
y,
sample_axis=[0, 2],
event_axis=[1],
keepdims=True)
self.assertAllEqual((1, 3, 3, 1, 5), cov_kd.shape)
cov_kd = self.evaluate(cov_kd)
self.assertAllClose(cov, cov_kd[0, :, :, 0, :])
for i in range(5): # Iterate over batch index.
# Get ith batch of samples, and permute/reshape to [n_samples, n_events]
x_i = np.reshape(np.transpose(x[:, :, :, i], [0, 2, 1]),
[2 * 4, 3])
y_i = np.reshape(np.transpose(y[:, :, :, i], [0, 2, 1]),
[2 * 4, 3])
# Will compare with ith batch of covariance.
cov_i = cov[:, :, i]
for m in range(3): # Iterate over row of matrix
for n in range(3): # Iterate over column of matrix
self.assertAllClose(self._np_cov_1d(x_i[:, m], y_i[:, n]),
cov_i[m, n])
def test_batch_vector_sampaxis13_eventaxis2(self):
rng = test_util.test_np_rng()
# x.shape = batch, sample, event, sample
x = rng.randn(4, 10, 2, 10)
y = x + 0.1 * rng.randn(10, 2, 10)
x[:, :, 0, :] += 0.1 * rng.randn(4, 10, 10)
cov = tfp.stats.covariance(x, y, sample_axis=[1, 3], event_axis=[2])
self.assertAllEqual((4, 2, 2), cov.shape)
cov = self.evaluate(cov)
cov_kd = tfp.stats.covariance(x,
y,
sample_axis=[1, 3],
event_axis=[2],
keepdims=True)
self.assertAllEqual((4, 1, 2, 2, 1), cov_kd.shape)
cov_kd = self.evaluate(cov_kd)
self.assertAllClose(cov, cov_kd[:, 0, :, :, 0])
for i in range(4): # Iterate over batch index.
# Get ith batch of samples, and permute/reshape to [n_samples, n_events]
x_i = np.reshape(np.transpose(x[i, :, :, :], [0, 2, 1]),
[10 * 10, 2])
y_i = np.reshape(np.transpose(y[i, :, :, :], [0, 2, 1]),
[10 * 10, 2])
# Will compare with ith batch of covariance.
cov_i = cov[i, :, :]
for m in range(2): # Iterate over row of matrix
for n in range(2): # Iterate over column of matrix
self.assertAllClose(self._np_cov_1d(x_i[:, m], y_i[:, n]),
cov_i[m, n])
def test_batch_vector_sampaxis0_eventaxisn1(self):
rng = test_util.test_np_rng()
# X and Y are correlated, albeit less so in the first component.
# They both are both 100 samples of 3-batch vectors in R^2.
x = rng.randn(100, 3, 2)
y = x + 0.1 * rng.randn(100, 3, 2)
x[:, :, 0] += 0.1 * rng.randn(100, 3)
corr = tfp.stats.correlation(x, y, event_axis=-1)
self.assertAllEqual((3, 2, 2), corr.shape)
corr = self.evaluate(corr)
corr_kd = tfp.stats.correlation(x, y, event_axis=-1, keepdims=True)
self.assertAllEqual((1, 3, 2, 2), corr_kd.shape)
corr_kd = self.evaluate(corr_kd)
self.assertAllClose(corr, corr_kd[0, ...])
for i in range(3): # Iterate over batch index.
x_i = x[:, i, :] # Pick out ith batch of samples.
y_i = y[:, i, :]
corr_i = corr[i, :, :]
for m in range(2): # Iterate over row of matrix
for n in range(2): # Iterate over column of matrix
self.assertAllClose(self._np_corr_1d(x_i[:, m], y_i[:, n]),
corr_i[m, n])
def _build_model_and_params(self,
num_timesteps,
param_batch_shape,
num_posterior_draws=10):
seed = test_util.test_seed_stream()
rng = test_util.test_np_rng(seed())
observed_time_series = self._build_tensor(
rng.randn(*(param_batch_shape + [num_timesteps])))
# Build an STS model with multiple components
day_of_week = tfp.sts.Seasonal(
num_seasons=7,
observed_time_series=observed_time_series,
name='day_of_week')
local_linear_trend = tfp.sts.LocalLinearTrend(
observed_time_series=observed_time_series,
name='local_linear_trend')
model = tfp.sts.Sum(components=[day_of_week, local_linear_trend],
observed_time_series=observed_time_series)
# Sample test params from the prior (faster than posterior samples).
param_samples = [
p.prior.sample([num_posterior_draws], seed=seed())
for p in model.parameters
]
return model, observed_time_series, param_samples
def test_batch_vector_sampaxis1_eventaxis2(self):
rng = test_util.test_np_rng()
# x.shape = [2, 5000, 2],
# 2-batch members, 5000 samples each, events in R^2.
x0 = rng.randn(5000, 2)
x1 = 2 * rng.randn(5000, 2)
x = np.stack((x0, x1), axis=0)
# chol.shape = [2 (batch), 2x2 (event x event)]
chol = tfp.stats.cholesky_covariance(x, sample_axis=1)
chol_kd = tfp.stats.cholesky_covariance(x,
sample_axis=1,
keepdims=True)
# Make sure static shape of keepdims works
self.assertAllEqual((2, 2, 2), chol.shape)
self.assertAllEqual((2, 1, 2, 2), chol_kd.shape)
chol, chol_kd = self.evaluate([chol, chol_kd])
# keepdims should not change the numbers in the result.
self.assertAllEqual(chol, np.squeeze(chol_kd, axis=1))
# Covariance is trivial since these are independent normals.
# Tolerance chosen to be 2x the lowest passing atol.
self.assertAllClose(np.eye(2), chol[0, ...], atol=0.06)
self.assertAllClose(2 * np.eye(2), chol[1, ...], atol=0.06)
def test_random_mean(self):
rng = test_util.test_np_rng()
x = rng.rand(100)
running_mean = tfp.experimental.stats.RunningMean.from_shape(shape=())
running_mean = consume(running_mean, x)
mean = self.evaluate(running_mean.mean)
self.assertAllClose(np.mean(x), mean, rtol=1e-6)
def setUp(self):
super(MutualInformationTest, self).setUp()
self.seed = tfp_test_util.test_seed()
self.scores = tfp_test_util.test_np_rng().normal(loc=1.0,
scale=2.0,
size=[13, 17])
batch_size = 1000
rho = 0.8
dim = 2
x, eps = tf.split(value=tf.random.normal(shape=(2 * batch_size, dim),
seed=self.seed),
num_or_size_splits=2,
axis=0)
mean = rho * x
stddev = tf.sqrt(1. - tf.square(rho))
y = mean + stddev * eps
conditional_dist = tfd.MultivariateNormalDiag(
mean, scale_identity_multiplier=stddev)
marginal_dist = tfd.MultivariateNormalDiag(tf.zeros(dim), tf.ones(dim))
# The conditional_scores has its shape [y_batch_dim, distibution_batch_dim]
# as the `lower_bound_info_nce` requires `scores[i, j] = f(x[i], y[j])
# = log p(x[i] | y[j])`.
self.conditional_scores = conditional_dist.log_prob(y[:,
tf.newaxis, :])
self.marginal_scores = marginal_dist.log_prob(y)[:, tf.newaxis]
self.optimal_critic = 1 + self.conditional_scores - self.marginal_scores
self.theoretical_mi = np.float32(-0.5 * np.log(1. - rho**2) * dim)
# Y is N-D standard normal distributed.
self.differential_entropy_y = 0.5 * np.log(2 * np.pi * np.e) * dim
def test_running_variance(self, ddof): rng = test_util.test_np_rng() x = rng.rand(100) var = tfp.experimental.stats.RunningVariance.from_shape() var = consume(var, x) final_mean, final_var = self.evaluate([var.mean, var.variance(ddof=ddof)]) self.assertNear(np.mean(x), final_mean, err=1e-6) self.assertNear(np.var(x, ddof=ddof), final_var, err=1e-6)
def test_chunking(self):
rng = test_util.test_np_rng()
x = rng.rand(100, 10, 5)
running_mean = tfp.experimental.stats.RunningMean.from_shape(
shape=(5,))
running_mean = consume(running_mean, x, chunk_axis=0)
mean = self.evaluate(running_mean.mean)
self.assertAllClose(np.mean(x.reshape(1000, 5), axis=0), mean, rtol=1e-6)
def test_manual_dtype(self):
rng = test_util.test_np_rng()
x = rng.rand(3, 10)
cov = tfp.experimental.stats.RunningCovariance.from_shape(
(10, ), dtype=tf.float64)
cov = consume(cov, x)
final_cov = cov.covariance()
self.assertEqual(final_cov.dtype, tf.float64)
def test_overlapping_axis_raises(self):
rng = test_util.test_np_rng()
# They both are both 100 samples of 3-batch vectors in R^2.
x = rng.randn(100, 3, 2)
y = x + 0.1 * rng.randn(100, 3, 2)
with self.assertRaisesRegexp(ValueError, 'overlapped'):
tfp.stats.covariance(x, y, sample_axis=[0, 1], event_axis=[1, 2])
def test_non_contiguous_event_axis_raises(self):
rng = test_util.test_np_rng()
# They both are both 100 samples of 3-batch vectors in R^2.
x = rng.randn(100, 3, 2)
y = x + 0.1 * rng.randn(100, 3, 2)
with self.assertRaisesRegexp(ValueError, 'must be contiguous'):
tfp.stats.covariance(x, y, sample_axis=1, event_axis=[0, 2])
def testPrecomputedWithMasking(self):
amplitude = np.array([1., 2.], np.float64)
length_scale = np.array([[.1], [.2], [.3]], np.float64)
observation_noise_variance = np.array([[1e-2], [1e-4], [1e-6]],
np.float64)
rng = test_util.test_np_rng()
observations_is_missing = np.array([
[False, True, False, True, False, True],
[False, False, False, False, False, False],
[True, True, False, False, True, True],
]).reshape((3, 1, 6))
observation_index_points = np.where(
observations_is_missing[..., np.newaxis], np.nan,
rng.uniform(-1., 1., (3, 1, 6, 2)).astype(np.float64))
observations = np.where(
observations_is_missing, np.nan,
rng.uniform(-1., 1., (3, 1, 6)).astype(np.float64))
index_points = rng.uniform(-1., 1., (5, 2)).astype(np.float64)
kernel = psd_kernels.ExponentiatedQuadratic(amplitude, length_scale)
gprm = tfd.GaussianProcessRegressionModel.precompute_regression_model(
kernel=kernel,
index_points=index_points,
observation_index_points=observation_index_points,
observations=observations,
observations_is_missing=observations_is_missing,
observation_noise_variance=observation_noise_variance,
validate_args=True)
self.assertAllNotNan(gprm.mean())
self.assertAllNotNan(gprm.variance())
self.assertAllNotNan(gprm.covariance())
# For each batch member of `gprm`, check that the distribution is the same
# as a GaussianProcessRegressionModel with no masking but conditioned on
# only the not-masked-out index points.
x = gprm.sample(seed=test_util.test_seed())
for i in range(3):
observation_index_points_i = tf.gather(
observation_index_points[i, 0],
(~observations_is_missing[i, 0]).nonzero()[0])
observations_i = tf.gather(
observations[i,
0], (~observations_is_missing[i, 0]).nonzero()[0])
gprm_i = tfd.GaussianProcessRegressionModel.precompute_regression_model(
kernel=kernel[i],
index_points=index_points,
observation_index_points=observation_index_points_i,
observations=observations_i,
observation_noise_variance=observation_noise_variance[i, 0],
validate_args=True)
self.assertAllClose(gprm.mean()[i], gprm_i.mean())
self.assertAllClose(gprm.variance()[i], gprm_i.variance())
self.assertAllClose(gprm.covariance()[i], gprm_i.covariance())
self.assertAllClose(gprm.log_prob(x)[i], gprm_i.log_prob(x[i]))
def test_sorted_descending_running_covariance(self):
rng = test_util.test_np_rng()
x = rng.rand(100, 10)
x[::-1].sort(axis=0) # sorts in descending order
cov = tfp.experimental.stats.RunningCovariance.from_shape((10, ))
for sample in x:
cov = cov.update(sample)
final_cov = self.evaluate(cov.covariance())
self.assertAllClose(final_cov, np.cov(x.T, ddof=0), rtol=1e-5)
def test_diagonal_of_correlation_matrix_x_with_x_is_one(self):
rng = test_util.test_np_rng()
# Some big numbers, to test stability.
x = np.float32(1e10 * rng.randn(100, 3))
corr = tfp.stats.correlation(x, sample_axis=0, event_axis=1)
self.assertAllEqual((3, 3), corr.shape)
corr = self.evaluate(corr)
self.assertAllClose([1., 1., 1.], np.diag(corr))
def test_random_mean(self):
rng = test_util.test_np_rng()
x = rng.rand(100)
running_mean = tfp.experimental.stats.RunningMean(shape=(), )
state = running_mean.initialize()
for sample in x:
state = running_mean.update(state, sample)
mean = self.evaluate(running_mean.finalize(state))
self.assertAllClose(np.mean(x), mean, rtol=1e-6)
def _random_batch_psd(self, dim):
rng = test_util.test_np_rng()
matrix = rng.random_sample([2, dim, dim])
matrix = np.matmul(matrix, np.swapaxes(matrix, -2, -1))
matrix = (matrix + np.diag(np.arange(dim) * .1)).astype(self.dtype)
masked_shape = (
matrix.shape if self.use_static_shape else [None] * len(matrix.shape))
matrix = tf1.placeholder_with_default(matrix, shape=masked_shape)
return matrix
def test_manual_dtype(self):
rng = test_util.test_np_rng()
x = rng.rand(100, 10)
cov = tfp.experimental.stats.RunningCovariance.from_shape(
(10, ), dtype=tf.float64)
for sample in x:
cov = cov.update(sample)
final_cov = cov.covariance()
self.assertEqual(final_cov.dtype, tf.float64)
def test_running_covariance(self, ddof):
rng = test_util.test_np_rng()
x = rng.rand(100, 10)
cov = tfp.experimental.stats.RunningCovariance.from_shape((10, ))
for sample in x:
cov = cov.update(sample)
final_mean, final_cov = self.evaluate(
[cov.mean, cov.covariance(ddof=ddof)])
self.assertAllClose(np.mean(x, axis=0), final_mean, rtol=1e-5)
self.assertAllClose(np.cov(x.T, ddof=ddof), final_cov, rtol=1e-5)
def test_running_covariance(self, ddof): rng = test_util.test_np_rng() x = rng.rand(100, 10) cov = tfp.experimental.stats.RunningCovariance.from_shape((10,)) cov = consume(cov, x) final_mean, final_cov = self.evaluate([cov.mean, cov.covariance(ddof=ddof)]) self.assertAllClose(np.mean(x, axis=0), final_mean, rtol=1e-5) self.assertAllClose(np.cov(x.T, ddof=ddof), final_cov, rtol=1e-5) self.assertEqual(cov.event_ndims, 1) self.assertEqual(cov.mean.dtype, tf.float32)
def test_axis_2_center_true_max_lags_1(self):
rng = test_util.test_np_rng()
x = rng.randn(3, 4, 5).astype(self.dtype)
if self.dtype in [np.complex64]:
x = 1j * rng.randn(3, 4, 5).astype(self.dtype)
self.check_results_versus_brute_force(x,
axis=2,
max_lags=1,
center=True,
normalize=False)
def test_axis_n1_center_false_max_lags_none_normalize_true(self):
rng = test_util.test_np_rng()
x = rng.randn(2, 3, 4).astype(self.dtype)
if self.dtype in [np.complex64]:
x = 1j * rng.randn(2, 3, 4).astype(self.dtype)
self.check_results_versus_brute_force(x,
axis=-1,
max_lags=None,
center=False,
normalize=True)
def test_manual_dtype(self):
rng = test_util.test_np_rng()
x = rng.rand(100, 10)
running_cov = tfp.experimental.stats.RunningCovariance(
(10, ), dtype=tf.float64)
state = running_cov.initialize()
for sample in x:
state = running_cov.update(state, sample)
final_cov = running_cov.finalize(state)
self.assertEqual(final_cov.dtype, tf.float64)
def test_iid_normal_passes(self):
n_samples = 500
# five scalar chains taken from iid Normal(0, 1)
rng = test_util.test_np_rng()
iid_normal_samples = rng.randn(n_samples, 5)
rhat_reducer = tfp.experimental.mcmc.PotentialScaleReductionReducer(
independent_chain_ndims=1)
rhat = self.evaluate(test_fixtures.reduce(rhat_reducer, iid_normal_samples))
self.assertAllEqual((), rhat.shape)
self.assertAllClose(1., rhat, rtol=0.02)
def test_sorted_ascending_running_covariance(self):
rng = test_util.test_np_rng()
x = rng.rand(100, 10)
x.sort(axis=0)
running_cov = tfp.experimental.stats.RunningCovariance((10, ))
state = running_cov.initialize()
for sample in x:
state = running_cov.update(state, sample)
final_cov = self.evaluate(running_cov.finalize(state))
self.assertAllClose(final_cov, np.cov(x.T, ddof=0), rtol=1e-5)
def test_random_higher_rank_samples(self):
rng = test_util.test_np_rng()
x_orig = rng.rand(100, 10)
x = tf.convert_to_tensor(x_orig, dtype=tf.float32)
running_moments = tfp.experimental.stats.RunningCentralMoments.from_shape(
shape=(10, ), moment=np.arange(5) + 1)
running_moments = consume(running_moments, x)
moments = self.evaluate(running_moments.moments())
self.assertAllClose(stats.moment(x_orig, moment=[1, 2, 3, 4, 5]),
moments,
rtol=1e-6)
