def test_posterior_mode_high_rank(self, rank_o, rank_t, rank_i, rank_s):
def increase_rank(n, x):
# By choosing prime number dimensions we make it less
# likely that a test will pass for accidental reasons.
primes = [3, 5, 7]
for i in range(n):
x = primes[i] * [x]
return x
observation_locs_data = tf.constant(increase_rank(
rank_o, [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]),
dtype=self.dtype)
observation_scales_data = tf.constant([0.25, 0.25, 0.25, 0.25],
dtype=self.dtype)
transition_matrix_data = tf.constant(increase_rank(
rank_t, [[0.8, 0.1, 0.1, 0.0], [0.1, 0.8, 0.0, 0.1],
[0.1, 0.0, 0.8, 0.1], [0.0, 0.1, 0.1, 0.8]]),
dtype=self.dtype)
initial_prob_data = tf.constant(increase_rank(
rank_i, [0.25, 0.25, 0.25, 0.25]),
dtype=self.dtype)
(initial_prob, transition_matrix, observation_locs,
observation_scales) = self.make_placeholders([
initial_prob_data, transition_matrix_data, observation_locs_data,
observation_scales_data
])
observations = tf.constant(increase_rank(
rank_s,
[[[0.91, 0.11], [0.21, 0.09]], [[0.11, 0.97], [0.12, 0.08]],
[[0.01, 0.12], [0.92, 0.11]], [[0.02, 0.11], [0.77, 0.11]],
[[0.81, 0.15], [0.21, 0.03]], [[0.01, 0.13], [0.23, 0.91]],
[[0.11, 0.12], [0.23, 0.79]], [[0.13, 0.11], [0.91, 0.29]]]),
dtype=self.dtype)
observation_distribution = tfp.distributions.TransformedDistribution(
tfd.MultivariateNormalDiag(observation_locs,
scale_diag=observation_scales),
tfp.bijectors.Reshape((2, 2)))
[num_steps] = self.make_placeholders([8])
model = tfd.HiddenMarkovModel(tfd.Categorical(probs=initial_prob),
tfd.Categorical(probs=transition_matrix),
observation_distribution,
num_steps=num_steps,
validate_args=True)
inferred_states = model.posterior_mode(observations)
rank_e = max(rank_o, rank_t, rank_i, rank_s)
expected_states = increase_rank(rank_e, [0, 2, 2, 2, 2, 3, 3, 2])
self.assertAllEqual(inferred_states, expected_states)
def test_posterior_mode_invariance_states(self):
observation_probs_data = tf.constant([[0.12, 0.48, 0.5, 0.1],
[0.4, 0.1, 0.5, 0.0],
[0.1, 0.2, 0.3, 0.4]],
dtype=self.dtype)
transition_matrix_data = tf.constant([[0.21, 0.49, 0.3],
[0.18, 0.12, 0.7],
[0.75, 0.15, 0.1]],
dtype=self.dtype)
initial_prob_data = tf.constant([0.8, 0.13, 0.07],
dtype=self.dtype)
(initial_prob, transition_matrix,
observation_probs) = self.make_placeholders([
initial_prob_data, transition_matrix_data,
observation_probs_data])
permutations = tf.identity(np.array([np.random.permutation(3)
for _ in range(8)]))
inverse_permutations = tf.argsort(permutations)
initial_prob_permuted = tf.gather(initial_prob, inverse_permutations)
# Permute rows of observation matrix
observation_probs_permuted = tf.gather(observation_probs,
inverse_permutations)
# Permute both rows and columns of transition matrix
transition_matrix_permuted = tf.transpose(
a=tf.gather(tf.transpose(a=transition_matrix),
inverse_permutations),
perm=[0, 2, 1])
transition_matrix_permuted = tf1.batch_gather(transition_matrix_permuted,
inverse_permutations)
observations = tf.constant([1, 0, 3, 1, 3, 0, 2, 1, 2, 1, 3, 0, 0, 1, 1, 2])
[num_steps] = self.make_placeholders([16])
model = tfd.HiddenMarkovModel(
tfd.Categorical(probs=initial_prob_permuted),
tfd.Categorical(probs=transition_matrix_permuted),
tfd.Categorical(probs=observation_probs_permuted),
num_steps=num_steps)
inferred_states = model.posterior_mode(observations)
expected_states = [0, 1, 2, 0, 2, 1, 2, 0, 2, 0, 2, 0, 1, 2, 0, 1]
expected_states_permuted = tf.transpose(
a=tf1.batch_gather(
tf.expand_dims(tf.transpose(
a=permutations), axis=-1), expected_states)[..., 0])
self.assertAllEqual(inferred_states, expected_states_permuted)
def distribution_fn(sample):
num_frames = sample.shape[-1]
mask = tf.one_hot(0, num_frames)[:, tf.newaxis]
probs = tf.roll(tf.one_hot(sample, 3), shift=1, axis=-2)
probs = probs * (1.0 - mask) + tf.convert_to_tensor([0.5, 0.5, 0]) * mask
return tfd.Independent(tfd.Categorical(probs=probs),
reinterpreted_batch_ndims=1)
def test_bug170030378(self):
n_item = 50
n_rater = 7
stream = test_util.test_seed_stream()
weight = self.evaluate(
tfd.Sample(tfd.Dirichlet([0.25, 0.25]), n_item).sample(seed=stream()))
mixture_dist = tfd.Categorical(probs=weight) # batch_shape=[50]
rater_sensitivity = self.evaluate(
tfd.Sample(tfd.Beta(5., 1.), n_rater).sample(seed=stream()))
rater_specificity = self.evaluate(
tfd.Sample(tfd.Beta(2., 5.), n_rater).sample(seed=stream()))
probs = tf.stack([rater_sensitivity, rater_specificity])[None, ...]
components_dist = tfd.BatchBroadcast( # batch_shape=[50, 2]
tfd.Independent(tfd.Bernoulli(probs=probs),
reinterpreted_batch_ndims=1),
[50, 2])
obs_dist = tfd.MixtureSameFamily(mixture_dist, components_dist)
observed = self.evaluate(obs_dist.sample(seed=stream()))
mixture_logp = obs_dist.log_prob(observed)
expected_logp = tf.math.reduce_logsumexp(
tf.math.log(weight) + components_dist.distribution.log_prob(
observed[:, None, ...]),
axis=-1)
self.assertAllClose(expected_logp, mixture_logp)
def testBrokenTypes(self):
with self.assertRaisesWithPredicateMatch(TypeError, 'Categorical'):
tfd.Mixture(None, [], validate_args=True)
cat = tfd.Categorical([0.3, 0.2])
# components must be a list of distributions
with self.assertRaisesWithPredicateMatch(
TypeError, 'all .* must be Distribution instances'):
tfd.Mixture(cat, [None], validate_args=True)
with self.assertRaisesWithPredicateMatch(TypeError, 'same dtype'):
tfd.Mixture(cat, [
tfd.Normal(loc=[1.0], scale=[2.0]),
tfd.Normal(loc=[np.float16(1.0)], scale=[np.float16(2.0)]),
],
validate_args=True)
with self.assertRaisesWithPredicateMatch(ValueError, 'non-empty list'):
tfd.Mixture(tfd.Categorical([0.3, 0.2]), None, validate_args=True)
def new(params,
event_size,
num_components,
dtype=None,
validate_args=False,
name=None):
"""Create the distribution instance from a `params` vector."""
with tf.name_scope(name, 'CategoricalMixtureOfOneHotCategorical',
[params, event_size, num_components]):
components_shape = tf.concat(
[tf.shape(params)[:-1], [num_components, event_size]], axis=0)
dist = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(
logits=params[..., :num_components],
validate_args=validate_args),
components_distribution=tfd.OneHotCategorical(
logits=tf.reshape(params[..., num_components:],
components_shape),
dtype=dtype or params.dtype.base_dtype,
validate_args=False
), # So we can eval on simplex interior.
# TODO(b/120154797): Change following to `validate_args=True` after
# fixing: "ValueError: `mixture_distribution` must have scalar
# `event_dim`s." assertion in MixtureSameFamily.
validate_args=False)
# pylint: disable=protected-access
dist._mean = functools.partial(_eval_all_one_hot,
tfd.Distribution.prob, dist)
dist.log_mean = functools.partial(_eval_all_one_hot,
tfd.Distribution.log_prob, dist)
# pylint: enable=protected-access
return dist
def testStatefulComponentDist(self):
class StatefulNormal(tfd.Distribution):
def __init__(self, loc):
self._loc = tf.convert_to_tensor(loc)
super(StatefulNormal, self).__init__(
dtype=tf.float32,
reparameterization_type=tfd.FULLY_REPARAMETERIZED,
validate_args=False,
allow_nan_stats=False)
def _batch_shape(self):
return self._loc.shape
def _event_shape(self):
return []
def _sample_n(self, n, seed=None):
return self._loc + tf.random.normal(
tf.concat([[n], tf.shape(self._loc)], axis=0), seed=seed)
mix = tfd.Mixture(
cat=tfd.Categorical(logits=[0., 0]),
components=[tfd.HalfNormal(scale=2.),
StatefulNormal(loc=3.)])
with warnings.catch_warnings(record=True) as triggered:
self.evaluate(mix.sample(seed=test_util.test_seed()))
self.assertTrue(
any('Falling back to stateful sampling for `components[1]`' in str(
warning.message) for warning in triggered))
def new(params,
num_components,
component_layer,
validate_args=False,
name=None,
**kwargs):
"""Create the distribution instance from a `params` vector."""
with tf.name_scope(name, 'MixtureSameFamily',
[params, num_components, component_layer]):
params = tf.convert_to_tensor(params, name='params')
num_components = tf.convert_to_tensor(num_components,
name='num_components',
preferred_dtype=tf.int32)
components_dist = component_layer(
tf.reshape(
params[..., num_components:],
tf.concat([tf.shape(params)[:-1], [num_components, -1]],
axis=0)))
mixture_dist = tfd.Categorical(logits=params[..., :num_components])
return tfd.MixtureSameFamily(
mixture_dist,
components_dist,
# TODO(b/120154797): Change following to `validate_args=True` after
# fixing: "ValueError: `mixture_distribution` must have scalar
# `event_dim`s." assertion in MixtureSameFamily.
validate_args=False,
**kwargs)
def testGradientsThroughParams(self):
logits = tf.Variable(np.zeros((3, 5, 2)),
dtype=tf.float32,
shape=tf.TensorShape([None, None, 2]))
concentration = tf.Variable(np.ones((3, 5, 4)),
dtype=tf.float32,
shape=tf.TensorShape(None))
loc = tf.Variable(np.zeros((3, 5, 4)),
dtype=tf.float32,
shape=tf.TensorShape(None))
scale = tf.Variable(1., dtype=tf.float32, shape=tf.TensorShape(None))
dist = tfd.Mixture(tfd.Categorical(logits=logits),
components=[
tfd.Dirichlet(concentration),
tfd.MultivariateNormalDiag(
loc=loc, scale_identity_multiplier=scale)
],
use_static_graph=self.use_static_graph,
validate_args=True)
with tf.GradientTape() as tape:
loss = tf.reduce_sum(dist.log_prob(tf.ones((3, 5, 4)) / 4.))
grad = tape.gradient(loss, dist.trainable_variables)
self.assertLen(grad, 4)
self.assertAllNotNone(grad)
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)
logits = tf.random.uniform(tf.concat(
(batch_shape_tensor, [num_components]), 0),
-1,
1,
dtype=tf.float32) - 50.
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)
scale_diag = 10 * tf.random.uniform(batch_and_event_shape)
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)
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)
def testCdfBatchUnivariate(self):
"""Tests against scipy for a (batch of) mixture(s) of seven gaussians."""
n_components = 7
batch_size = 5
mixture_weight_logits = np.random.uniform(
low=-1, high=1, size=(batch_size, n_components)).astype(np.float32)
def _batch_univariate_softmax(x):
e_x = np.exp(x)
e_x_sum = np.expand_dims(np.sum(e_x, axis=1), axis=1)
return e_x / np.tile(e_x_sum, reps=[1, x.shape[1]])
psize = (batch_size, )
mixture_weights = _batch_univariate_softmax(mixture_weight_logits)
means = [
np.random.uniform(low=-10, high=10, size=psize).astype(np.float32)
for _ in range(n_components)
]
sigmas = [
np.ones(shape=psize, dtype=np.float32) for _ in range(n_components)
]
cat_tf = tfd.Categorical(probs=mixture_weights)
components_tf = [
tfd.Normal(loc=mu, scale=sigma)
for (mu, sigma) in zip(means, sigmas)
]
mixture_tf = tfd.Mixture(cat=cat_tf,
components=components_tf,
use_static_graph=self.use_static_graph,
validate_args=True)
xs_to_check = [
np.array([1.0, 5.9, -3, 0.0, 0.0], dtype=np.float32),
np.random.randn(batch_size).astype(np.float32)
]
for x_tensor in xs_to_check:
with self.cached_session() as sess:
x_cdf_tf_result, x_log_cdf_tf_result = sess.run(
[mixture_tf.cdf(x_tensor),
mixture_tf.log_cdf(x_tensor)])
# Compute the cdf with scipy.
scipy_component_cdfs = [
stats.norm.cdf(x=x_tensor, loc=mu, scale=sigma)
for (mu, sigma) in zip(means, sigmas)
]
weights_and_cdfs = zip(
np.transpose(mixture_weights, axes=[1, 0]),
scipy_component_cdfs)
final_cdf_probs_per_component = [
np.multiply(c_p_value, d_cdf_value)
for (c_p_value, d_cdf_value) in weights_and_cdfs
]
scipy_cdf_result = np.sum(final_cdf_probs_per_component,
axis=0)
self.assertAllClose(x_cdf_tf_result, scipy_cdf_result)
self.assertAllClose(np.exp(x_log_cdf_tf_result),
scipy_cdf_result)
class BatchShapeInferenceTests(test_util.TestCase):
@parameterized.named_parameters(
{
'testcase_name': '_trivial',
'value_fn': lambda: tfd.Normal(loc=0., scale=1.),
'expected_batch_shape': []
},
{
'testcase_name':
'_simple_tensor_broadcasting',
'value_fn':
lambda: tfd.MultivariateNormalDiag( # pylint: disable=g-long-lambda
loc=[0., 0.],
scale_diag=tf.convert_to_tensor([[1., 1.], [1., 1.]])),
'expected_batch_shape': [2]
},
{
'testcase_name':
'_rank_deficient_tensor_broadcasting',
'value_fn':
lambda: tfd.MultivariateNormalDiag( # pylint: disable=g-long-lambda
loc=0.,
scale_diag=tf.convert_to_tensor([[1., 1.], [1., 1.]])),
'expected_batch_shape': [2]
},
{
'testcase_name':
'_mixture_same_family',
'value_fn':
lambda: tfd.MixtureSameFamily( # pylint: disable=g-long-lambda
mixture_distribution=tfd.Categorical(logits=[[[1., 2., 3.],
[4., 5., 6.]]]),
components_distribution=tfd.Normal(
loc=0., scale=[[[1., 2., 3.], [4., 5., 6.]]])),
'expected_batch_shape': [1, 2]
},
{
'testcase_name':
'_deeply_nested',
'value_fn':
lambda: tfd.Independent( # pylint: disable=g-long-lambda
tfd.Independent(tfd.Independent(tfd.Independent(
tfd.Normal(loc=0., scale=[[[[[[[[1.]]]]]]]]),
reinterpreted_batch_ndims=2),
reinterpreted_batch_ndims=0),
reinterpreted_batch_ndims=1),
reinterpreted_batch_ndims=1),
'expected_batch_shape': [1, 1, 1, 1]
})
def test_batch_shape_inference_is_correct(self, value_fn,
expected_batch_shape):
value = value_fn(
) # Defer construction until we're in the right graph.
self.assertAllEqual(expected_batch_shape,
value._inferred_batch_shape_tensor())
batch_shape = value._inferred_batch_shape()
self.assertIsInstance(batch_shape, tf.TensorShape)
self.assertTrue(batch_shape.is_compatible_with(expected_batch_shape))
def testExcessiveConcretizationOfParams(self):
logits = tfp_hps.defer_and_count_usage(
tf.Variable(np.zeros((3, 5, 2)),
dtype=tf.float32,
shape=tf.TensorShape([None, None, 2]),
name='logits'))
concentration = tfp_hps.defer_and_count_usage(
tf.Variable(np.ones((3, 5, 4)),
dtype=tf.float32,
shape=tf.TensorShape(None),
name='concentration'))
loc = tfp_hps.defer_and_count_usage(
tf.Variable(np.zeros((3, 5, 4)),
dtype=tf.float32,
shape=tf.TensorShape(None),
name='loc'))
scale = tfp_hps.defer_and_count_usage(
tf.Variable(1.,
dtype=tf.float32,
shape=tf.TensorShape(None),
name='scale'))
dist = tfd.Mixture(tfd.Categorical(logits=logits),
components=[
tfd.Dirichlet(concentration),
tfd.Independent(tfd.Normal(loc=loc,
scale=scale),
reinterpreted_batch_ndims=1)
],
use_static_graph=self.use_static_graph,
validate_args=True)
for method in ('batch_shape_tensor', 'event_shape_tensor',
'entropy_lower_bound'):
with tfp_hps.assert_no_excessive_var_usage(method,
max_permissible=2):
getattr(dist, method)()
with tfp_hps.assert_no_excessive_var_usage('sample',
max_permissible=2):
dist.sample(seed=test_util.test_seed())
for method in ('prob', 'log_prob'):
with tfp_hps.assert_no_excessive_var_usage('method',
max_permissible=2):
getattr(dist, method)(tf.ones((3, 5, 4)) / 4.)
# TODO(b/140579567): The `stddev()` and `variance()` methods require
# calling both:
# - `self.components[i].mean()`
# - `self.components[i].stddev()`
# Thus, these methods incur an additional concretization (or two if
# `validate_args=True` for `self.components[i]`).
for method in ('stddev', 'variance'):
with tfp_hps.assert_no_excessive_var_usage(method,
max_permissible=3):
getattr(dist, method)()
def testBrokenShapeUnknownCategories(self):
with self.assertRaisesWithPredicateMatch(ValueError,
r'Could not infer'):
cat_logits = tf.Variable([[13., 19.]],
shape=[1, None],
dtype=tf.float32)
tfd.Mixture(tfd.Categorical(cat_logits),
[tfd.Normal(loc=[1.0], scale=[2.0])],
validate_args=True)
def testBrokenTypes(self):
with self.assertRaisesWithPredicateMatch(TypeError, "Categorical"):
tfd.Mixture(None, [], use_static_graph=self.use_static_graph)
cat = tfd.Categorical([0.3, 0.2])
# components must be a list of distributions
with self.assertRaisesWithPredicateMatch(
TypeError, "all .* must be Distribution instances"):
tfd.Mixture(cat, [None], use_static_graph=self.use_static_graph)
with self.assertRaisesWithPredicateMatch(TypeError, "same dtype"):
tfd.Mixture(cat, [
tfd.Normal(loc=[1.0], scale=[2.0]),
tfd.Normal(loc=[np.float16(1.0)], scale=[np.float16(2.0)]),
],
use_static_graph=self.use_static_graph)
with self.assertRaisesWithPredicateMatch(ValueError, "non-empty list"):
tfd.Mixture(tfd.Categorical([0.3, 0.2]),
None,
use_static_graph=self.use_static_graph)
def test_single_sequence_posterior_marginals(self):
# In this test we have a 9-vertex graph with precisely one
# 7-vertex path from vertex 0 to vertex 8.
# The hidden Markov model is a random walk on this
# graph where the only observations are
# "observed at 0", "observed in {1, 2, ..., 7}",
# "observed at 8".
# The purpose of this test is to ensure that transition
# and observation matrices with many log probabilities
# equal to -infinity, and where the result contains many
# -infinities, are handled correctly.
initial_prob = tf.constant(np.ones(9) / 9.0, dtype=self.dtype)
edges = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 6), (2, 5), (5, 6),
(6, 7), (6, 8)]
transition_matrix = np.zeros((9, 9))
for (i, j) in edges:
transition_matrix[i, j] = 1.
transition_matrix[j, i] = 1.
transition_matrix = tf.constant(
transition_matrix /
np.sum(transition_matrix, axis=1, keepdims=True),
dtype=self.dtype)
observation_probs = tf.constant(np.block(
[[1, 0, 0], [np.zeros((7, 1)),
np.ones((7, 1)),
np.zeros((7, 1))], [0, 0, 1]]),
dtype=self.dtype)
[num_steps] = self.make_placeholders([7])
model = tfd.HiddenMarkovModel(tfd.Categorical(probs=initial_prob),
tfd.Categorical(probs=transition_matrix),
tfd.Categorical(probs=observation_probs),
num_steps=num_steps,
validate_args=True)
observations = [0, 1, 1, 1, 1, 1, 2]
probs = self.evaluate(
model.posterior_marginals(observations).probs_parameter())
expected_probs = np.eye(9)[[0, 1, 2, 3, 4, 6, 8]]
self.assertAllClose(probs, expected_probs, rtol=1e-4, atol=0.0)
def test_hmm(self):
n_steps = 4
infer_step = 2
observations = [-1.0, 0.0, 1.0, 2.0]
initial_prob = tf.constant([0.6, 0.4], dtype=tf.float32)
transition_matrix = tf.constant([[0.6, 0.4], [0.3, 0.7]],
dtype=tf.float32)
observation_locs = tf.constant([0.0, 1.0], dtype=tf.float32)
observation_scale = tf.constant(0.5, dtype=tf.float32)
dist1 = tfd.HiddenMarkovModel(tfd.Categorical(probs=initial_prob),
tfd.Categorical(probs=transition_matrix),
tfd.Normal(loc=observation_locs,
scale=observation_scale),
num_steps=n_steps)
p = dist1.posterior_marginals(
observations).probs_parameter()[infer_step]
def model():
i = yield Root(tfd.Categorical(probs=initial_prob, dtype=tf.int32))
z = yield tfd.Normal(loc=tf.gather(observation_locs, i),
scale=observation_scale)
for t in range(n_steps - 1):
i = yield tfd.Categorical(probs=tf.gather(
transition_matrix, i),
dtype=tf.int32)
yield tfd.Normal(loc=tf.gather(observation_locs, i),
scale=observation_scale)
dist2 = marginalize.MarginalizableJointDistributionCoroutine(model)
full_observations = list(
itertools.chain(*zip([
'tabulate' if i == infer_step else 'marginalize'
for i in range(n_steps)
], observations)))
q = tf.exp(dist2.marginalized_log_prob(full_observations))
q = q / tf.reduce_sum(q)
self.assertAllClose(p, q)
def testBrokenShapesStatic(self):
with self.assertRaisesWithPredicateMatch(ValueError,
r'cat.num_classes != len'):
tfd.Mixture(
tfd.Categorical([0.1, 0.5]), # 2 classes
[tfd.Normal(loc=1.0, scale=2.0)],
validate_args=True)
with self.assertRaisesWithPredicateMatch(
ValueError, r'components\[1\] batch shape must be compatible'):
# The value error is raised because the batch shapes of the
# Normals are not equal. One is a scalar, the other is a
# vector of size (2,).
tfd.Mixture(
tfd.Categorical([-0.5, 0.5]), # scalar batch
[
tfd.Normal(loc=1.0, scale=2.0), # scalar dist
tfd.Normal(loc=[1.0, 1.0], scale=[2.0, 2.0])
],
validate_args=True)
def test_posterior_mode_missing_discrete_observations(self):
initial_prob = tf.constant([1.0, 0.0, 0.0, 0.0], dtype=self.dtype)
# This test uses a model with a random walk that can make a change of
# of -1, 0 or +1 at each step.
transition_data = (0.5 * np.diag(np.ones(4)) +
0.25*np.diag(np.ones(3), -1) +
0.25*np.diag(np.ones(3), 1))
transition_data[0, 0] += 0.25
transition_data[3, 3] += 0.25
transition_matrix = tf.constant(transition_data, dtype=self.dtype)
# Observations of the random walk are unreliable and give the
# correct position with probability `0.25 + 0.75 * reliability`
def observation_fn(reliability):
return np.array(reliability * np.diag(np.ones(4)) +
(1 - reliability) * 0.25 * np.ones((4, 4)))
observation_data = np.array(
[observation_fn(reliability)
for reliability in [0.993, 0.994, 0.995, 0.996]])
observation_probs = tf.constant(observation_data, dtype=self.dtype)
[num_steps] = self.make_placeholders([7])
model = tfd.HiddenMarkovModel(tfd.Categorical(probs=initial_prob),
tfd.Categorical(probs=transition_matrix),
tfd.Categorical(probs=observation_probs),
num_steps=num_steps)
observations = tf.constant([0, 1, 2, 3, 2, 1, 0])
mask = tf.constant([False, True, True, False, True, True, False])
inferred_states = model.posterior_mode(observations, mask)
# This example has been tuned so that there are two local maxima in the
# space of paths.
# As the `reliability` parameter increases, the mode switches from one of
# the two paths to the other.
expected_states = [[0, 0, 0, 0, 0, 0, 0],
[0, 0, 0, 0, 0, 0, 0],
[0, 1, 2, 3, 2, 1, 0],
[0, 1, 2, 3, 2, 1, 0]]
self.assertAllEqual(inferred_states, expected_states)
def testCdfScalarUnivariate(self):
"""Tests CDF against scipy for a mixture of seven gaussians."""
# Construct a mixture of gaussians with seven components.
n_components = 7
# pre-softmax mixture probabilities.
mixture_weight_logits = np.random.uniform(
low=-1, high=1, size=(n_components, )).astype(np.float32)
def _scalar_univariate_softmax(x):
e_x = np.exp(x - np.max(x))
return e_x / e_x.sum()
# Construct the tfd.Mixture object.
mixture_weights = _scalar_univariate_softmax(mixture_weight_logits)
means = [
np.random.uniform(low=-10, high=10, size=()).astype(np.float32)
for _ in range(n_components)
]
sigmas = [
np.ones(shape=(), dtype=np.float32) for _ in range(n_components)
]
cat_tf = tfd.Categorical(probs=mixture_weights)
components_tf = [
tfd.Normal(loc=mu, scale=sigma)
for (mu, sigma) in zip(means, sigmas)
]
mixture_tf = tfd.Mixture(cat=cat_tf,
components=components_tf,
use_static_graph=self.use_static_graph,
validate_args=True)
# These are two test cases to verify.
xs_to_check = [
np.array(1.0, dtype=np.float32),
np.array(np.random.randn()).astype(np.float32)
]
for x_tensor in xs_to_check:
with self.cached_session() as sess:
x_cdf_tf_result, x_log_cdf_tf_result = sess.run(
[mixture_tf.cdf(x_tensor),
mixture_tf.log_cdf(x_tensor)])
# Compute the cdf with scipy.
scipy_component_cdfs = [
stats.norm.cdf(x=x_tensor, loc=mu, scale=sigma)
for (mu, sigma) in zip(means, sigmas)
]
scipy_cdf_result = np.dot(mixture_weights,
np.array(scipy_component_cdfs))
self.assertAllClose(x_cdf_tf_result, scipy_cdf_result)
self.assertAllClose(np.exp(x_log_cdf_tf_result),
scipy_cdf_result)
def test_coin_tosses(self):
initial_prob_data = tf.constant([0.5, 0.5], dtype=self.dtype)
transition_matrix_data = tf.constant([[0.5, 0.5],
[0.5, 0.5]], dtype=self.dtype)
observation_probs_data = tf.constant([[1.0, 0.0],
[0.0, 1.0]], dtype=self.dtype)
(initial_prob, transition_matrix,
observation_probs) = self.make_placeholders([
initial_prob_data, transition_matrix_data,
observation_probs_data])
[num_steps] = self.make_placeholders([5])
model = tfd.HiddenMarkovModel(tfd.Categorical(probs=initial_prob),
tfd.Categorical(probs=transition_matrix),
tfd.Categorical(probs=observation_probs),
num_steps=num_steps)
x = model.log_prob([0, 0, 0, 0, 0])
self.assertAllClose(x, np.log(.5**5), rtol=1e-5, atol=0.0)
def test_docstring_example(self):
stream = test_util.test_seed_stream()
loc = tfp.random.spherical_uniform([10], 3, seed=stream())
components_dist = tfd.VonMisesFisher(mean_direction=loc, concentration=50.)
mixture_dist = tfd.Categorical(
logits=tf.random.uniform([500, 10], seed=stream()))
obs_dist = tfd.MixtureSameFamily(
mixture_dist, tfd.BatchBroadcast(components_dist, [500, 10]))
test_sites = tfp.random.spherical_uniform([20], 3, seed=stream())
lp = tfd.Sample(obs_dist, 20).log_prob(test_sites)
self.assertEqual([500], lp.shape)
self.evaluate(lp)
def testSampleBimixGamma(self):
"""Tests a bug in the underlying tfd.Gamma op.
Mixture's use of dynamic partition requires `random_gamma` correctly returns
an empty `Tensor`.
"""
gm = tfd.Mixture(cat=tfd.Categorical(probs=[.3, .7]),
components=[tfd.Gamma(1., 2.),
tfd.Gamma(2., 1.)],
use_static_graph=self.use_static_graph)
x_ = self.evaluate(gm.sample())
self.assertAllEqual([], x_.shape)
def test_posterior_mode_invariance_observations(self):
observation_probs_data = tf.constant([[0.09, 0.48, 0.52, 0.11],
[0.31, 0.21, 0.21, 0.27]],
dtype=self.dtype)
transition_matrix_data = tf.constant([[0.90, 0.10],
[0.30, 0.70]],
dtype=self.dtype)
initial_prob_data = tf.constant([0.65, 0.35],
dtype=self.dtype)
(initial_prob, transition_matrix,
observation_probs) = self.make_placeholders([
initial_prob_data, transition_matrix_data,
observation_probs_data])
permutations = tf.identity(np.array([np.random.permutation(4)
for _ in range(8)]))
inverse_permutations = tf.argsort(permutations)
observation_probs_permuted = tf.transpose(
a=tf.gather(tf.transpose(a=observation_probs),
inverse_permutations),
perm=[0, 2, 1])
observations = tf.constant([1, 0, 3, 1, 3, 0, 2, 1, 2, 1, 3, 0, 0, 1, 1, 2])
observations_permuted = tf.transpose(
a=tf.gather(tf.transpose(a=permutations)[..., tf.newaxis],
observations,
batch_dims=observations.shape.ndims-1)[..., 0])
[num_steps] = self.make_placeholders([16])
model = tfd.HiddenMarkovModel(
tfd.Categorical(probs=initial_prob),
tfd.Categorical(probs=transition_matrix),
tfd.Categorical(probs=observation_probs_permuted),
num_steps=num_steps)
inferred_states = model.posterior_mode(observations_permuted)
expected_states = [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 0, 0, 0]
self.assertAllEqual(inferred_states, 8*[expected_states])
def make_univariate_mixture(batch_shape, num_components):
batch_shape = tf.convert_to_tensor(value=batch_shape, dtype=tf.int32)
seed_stream = test_util.test_seed_stream('univariate_mixture')
logits = -50. + tf.random.uniform(
tf.concat((batch_shape, [num_components]), axis=0),
-1, 1, dtype=tf.float32, seed=seed_stream())
components = [
tfd.Normal(loc=tf.random.normal(batch_shape, seed=seed_stream()),
scale=10 * tf.random.uniform(batch_shape, seed=seed_stream()))
for _ in range(num_components)
]
cat = tfd.Categorical(logits, dtype=tf.int32)
return tfd.Mixture(cat, components, validate_args=True)
def test_posterior_marginals_edge_case_no_transitions(self):
# Test all eight combinations of a single state that is
# 1. unmasked/masked
# 2. observed at state 0/state 1
# 3. more likely started at state 0/state 1
initial_prob_data = tf.constant([[0.9, 0.1], [0.1, 0.9]], dtype=self.dtype)
transition_matrix_data = tf.constant([[0.5, 0.5],
[0.5, 0.5]], dtype=self.dtype)
observation_probs_data = tf.constant([[1.0, 0.0],
[0.0, 1.0]], dtype=self.dtype)
(initial_prob, transition_matrix,
observation_probs) = self.make_placeholders([
initial_prob_data, transition_matrix_data,
observation_probs_data])
[num_steps] = self.make_placeholders([1])
model = tfd.HiddenMarkovModel(tfd.Categorical(probs=initial_prob),
tfd.Categorical(probs=transition_matrix),
tfd.Categorical(probs=observation_probs),
num_steps=num_steps)
inferred_marginals = self.evaluate(
model.posterior_marginals(
observations=[[[0]], [[1]]],
mask=[[[[True]]], [[[False]]]]).probs_parameter())
# Result is a [2,2,2] batch of sequences of length 1 of
# [2]-vectors of probabilities.
expected_marginals = [[[[[0.9, 0.1]],
[[0.1, 0.9]]],
[[[0.9, 0.1]],
[[0.1, 0.9]]]],
[[[[1., 0.]],
[[1., 0.]]],
[[[0., 1.]],
[[0., 1.]]]]]
self.assertAllClose(inferred_marginals, expected_marginals)
def test_edge_case_sample_n_no_transitions(self):
initial_prob_data = tf.constant([0.5, 0.5], dtype=self.dtype)
transition_matrix_data = tf.constant([[0.5, 0.5],
[0.5, 0.5]], dtype=self.dtype)
observation_probs_data = tf.constant([[1.0, 0.0],
[0.0, 1.0]], dtype=self.dtype)
(initial_prob, transition_matrix,
observation_probs) = self.make_placeholders([
initial_prob_data, transition_matrix_data,
observation_probs_data])
[num_steps] = self.make_placeholders([1])
model = tfd.HiddenMarkovModel(tfd.Categorical(probs=initial_prob),
tfd.Categorical(probs=transition_matrix),
tfd.Categorical(probs=observation_probs),
num_steps=num_steps)
x = model._sample_n(1)
x_shape = self.evaluate(tf.shape(input=x))
self.assertAllEqual(x_shape, [1, 1])
def test_mean_and_variance(self):
initial_prob_data = tf.constant([0.6, 0.1, 0.3], dtype=self.dtype)
transition_matrix_data = tf.constant([[0.2, 0.6, 0.2],
[0.5, 0.3, 0.2],
[0.0, 1.0, 0.0]], dtype=self.dtype)
observation_locs_data = tf.constant([0.0, 1.0, 2.0], dtype=self.dtype)
observation_scale_data = tf.constant(0.5, dtype=self.dtype)
(initial_prob, transition_matrix,
observation_locs, observation_scale) = self.make_placeholders([
initial_prob_data, transition_matrix_data,
observation_locs_data, observation_scale_data])
[num_steps] = self.make_placeholders([5])
model = tfd.HiddenMarkovModel(tfd.Categorical(probs=initial_prob),
tfd.Categorical(probs=transition_matrix),
tfd.Normal(loc=observation_locs,
scale=observation_scale),
num_steps=num_steps)
self.run_test_sample_consistent_mean_variance(self.evaluate, model,
num_samples=100000,
rtol=0.03)
def new(params,
probs_input=False,
dtype=None,
validate_args=False,
name='CategoricalLayer'):
"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
return tfd.Categorical(
logits=params if not probs_input else None,
probs=tf.clip_by_value(params, 1e-8, 1 - 1e-8) \
if probs_input else None,
dtype=dtype or params.dtype,
validate_args=validate_args,
name=name)
def make_univariate_mixture(batch_shape, num_components, use_static_graph):
batch_shape = tf.convert_to_tensor(value=batch_shape, dtype=tf.int32)
logits = tf.random.uniform(tf.concat(
(batch_shape, [num_components]), axis=0),
-1,
1,
dtype=tf.float32) - 50.
components = [
tfd.Normal(loc=tf.random.normal(batch_shape),
scale=10 * tf.random.uniform(batch_shape))
for _ in range(num_components)
]
cat = tfd.Categorical(logits, dtype=tf.int32)
return tfd.Mixture(cat, components, use_static_graph=use_static_graph)
