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)