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 _build_sts(self, observed_time_series=None): max_timesteps = 100 num_features = 3 prior = tfd.Sample(tfd.Laplace(0., 1.), sample_shape=[num_features]) # LinearRegression components don't currently take an `observed_time_series` # argument, so they can't infer a prior batch shape. This means we have to # manually set the batch shape expected by the tests. dtype = np.float32 if observed_time_series is not None: observed_time_series_tensor, _ = ( sts_util.canonicalize_observed_time_series_with_mask( observed_time_series)) batch_shape = tf.shape(observed_time_series_tensor)[:-2] dtype = dtype_util.as_numpy_dtype( observed_time_series_tensor.dtype) prior = tfd.Sample(tfd.Laplace(tf.zeros(batch_shape, dtype=dtype), 1.), sample_shape=[num_features]) regression = LinearRegression(design_matrix=np.random.randn( max_timesteps, num_features).astype(dtype), weights_prior=prior) return Sum(components=[regression], observed_time_series=observed_time_series)
def test_transformed_affine(self): sample_shape = 3 mvn = tfd.Independent(tfd.Normal(loc=[0., 0], scale=1), 1) aff = tfb.Affine(scale_tril=[[0.75, 0.], [0.05, 0.5]]) def expected_lp(y): x = aff.inverse(y) # Ie, tf.random.normal([4, 3, 2]) fldj = aff.forward_log_det_jacobian(x, event_ndims=1) return tf.reduce_sum(mvn.log_prob(x) - fldj, axis=1) # Transform a Sample. d = tfd.TransformedDistribution(tfd.Sample(mvn, sample_shape, validate_args=True), bijector=aff) y = d.sample(4, seed=test_util.test_seed()) actual_lp = d.log_prob(y) self.assertAllEqual((4, ) + (sample_shape, ) + (2, ), y.shape) self.assertAllEqual((4, ), actual_lp.shape) self.assertAllClose(*self.evaluate([expected_lp(y), actual_lp]), atol=0., rtol=1e-3) # Sample a Transform. d = tfd.Sample(tfd.TransformedDistribution(mvn, bijector=aff), sample_shape, validate_args=True) y = d.sample(4, seed=test_util.test_seed()) actual_lp = d.log_prob(y) self.assertAllEqual((4, ) + (sample_shape, ) + (2, ), y.shape) self.assertAllEqual((4, ), actual_lp.shape) self.assertAllClose(*self.evaluate([expected_lp(y), actual_lp]), atol=0., rtol=1e-3)
def model(): i = yield Root( tfd.Sample(tfd.Categorical(probs=probs, dtype=tf.int32), sample_shape=[2])) # Note use of scalar `sample_shape` to test expansion of shape # to vector. j = yield tfd.Sample(tfd.Categorical(probs=probs, dtype=tf.int32), sample_shape=2)
def test_kl_divergence(self): q_scale = 2. p = tfd.Sample( tfd.Independent(tfd.Normal(loc=tf.zeros([3, 2]), scale=1), 1), [5, 4], validate_args=True) q = tfd.Sample( tfd.Independent(tfd.Normal(loc=tf.zeros([3, 2]), scale=2.), 1), [5, 4], validate_args=True) actual_kl = tfd.kl_divergence(p, q) expected_kl = ((5 * 4) * (0.5 * q_scale**-2. - 0.5 + np.log(q_scale)) * # Actual KL. np.ones([3]) * 2) # Batch, events. self.assertAllClose(expected_kl, self.evaluate(actual_kl))
def test_broadcast_event(self): d = tfd.Sample(tfd.Normal(0, 1, validate_args=True), 4, validate_args=True) # Batch of 2 events: works. two_batch = d.log_prob(tf.zeros([2, 4])) self.assertEqual((2,), self.evaluate(two_batch).shape) # Broadcast of event: works. self.assertAllEqual(two_batch, d.log_prob(tf.zeros([2, 1])))
def test_entropy(self): sample_shape = [3, 4] mvn = tfd.Independent(tfd.Normal(loc=0, scale=[[0.25, 0.5]]), 1) d = tfd.Sample(mvn, sample_shape, validate_args=True) expected_entropy = 12 * tf.reduce_sum(mvn.distribution.entropy(), axis=-1) actual_entropy = d.entropy() self.assertAllEqual(*self.evaluate([expected_entropy, actual_entropy]))
def test_legacy_dists(self): class StatefulNormal(tfd.Normal): def _sample_n(self, n, seed=None): return self.loc + self.scale * tf.random.normal( tf.concat([[n], self.batch_shape_tensor()], axis=0), seed=seed) # pylint: disable=bad-whitespace d = tfd.JointDistributionNamed(dict( e = tfd.Independent(tfd.Exponential(rate=[100, 120]), 1), loc = StatefulNormal(loc=0, scale=2.), scale=lambda e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]), m = tfd.Normal, x =lambda m: tfd.Sample(tfd.Bernoulli(logits=m), 12)), validate_args=True) # pylint: enable=bad-whitespace warnings.simplefilter('always') with warnings.catch_warnings(record=True) as w: d.sample(seed=test_util.test_seed()) self.assertRegexpMatches( str(w[0].message), r'Falling back to stateful sampling for.*of type.*StatefulNormal.*' r'component name "loc" and `dist.name` "Normal"', msg=w)
def testCompareToBijector(self): """Demonstrates equivalence between TD, Bijector approach and AR dist.""" sample_shape = np.int32([4, 5]) batch_shape = np.int32([]) event_size = np.int32(2) batch_event_shape = np.concatenate([batch_shape, [event_size]], axis=0) sample0 = tf.zeros(batch_event_shape) affine = tfb.ScaleMatvecTriL( scale_tril=self._random_scale_tril(event_size), validate_args=True) ar = tfd.Autoregressive(self._normal_fn(affine), sample0, validate_args=True) ar_flow = tfb.MaskedAutoregressiveFlow( is_constant_jacobian=True, shift_and_log_scale_fn=lambda x: [None, affine.forward(x)], validate_args=True) td = tfd.TransformedDistribution( # TODO(b/137665504): Use batch-adding meta-distribution to set the batch # shape instead of tf.zeros. distribution=tfd.Sample(tfd.Normal(tf.zeros(batch_shape), 1.), [event_size]), bijector=ar_flow, validate_args=True) x_shape = np.concatenate([sample_shape, batch_shape, [event_size]], axis=0) x = 2. * self._rng.random_sample(x_shape).astype(np.float32) - 1. td_log_prob_, ar_log_prob_ = self.evaluate( [td.log_prob(x), ar.log_prob(x)]) self.assertAllClose(td_log_prob_, ar_log_prob_, atol=0., rtol=1e-6)
def test_doc_string(self): # Generate data. n = 2000 x2 = np.random.randn(n).astype(dtype=np.float32) * 2. x1 = np.random.randn(n).astype(dtype=np.float32) + (x2 * x2 / 4.) data = np.stack([x1, x2], axis=-1) # Density estimation with MADE. made = tfb.AutoregressiveNetwork(params=2, hidden_units=[10, 10]) distribution = tfd.TransformedDistribution( distribution=tfd.Sample(tfd.Normal(0., 1.), [2]), bijector=tfb.MaskedAutoregressiveFlow(made)) # Construct and fit model. x_ = tfkl.Input(shape=(2,), dtype=tf.float32) log_prob_ = distribution.log_prob(x_) model = tfk.Model(x_, log_prob_) model.compile(optimizer=tf1.train.AdamOptimizer(), loss=lambda _, log_prob: -log_prob) batch_size = 25 model.fit(x=data, y=np.zeros((n, 0), dtype=np.float32), batch_size=batch_size, epochs=1, steps_per_epoch=1, # Usually `n // batch_size`. shuffle=True, verbose=True) # Use the fitted distribution. self.assertAllEqual((3, 1, 2), distribution.sample((3, 1)).shape) self.assertAllEqual( (3,), distribution.log_prob(np.ones((3, 2), dtype=np.float32)).shape)
def test_misshapen_event(self): d = tfd.Sample(tfd.Normal(0, 1, validate_args=True), 4, validate_args=True) with self.assertRaisesRegexp(ValueError, r'Incompatible shapes for broadcasting'): self.evaluate(d.log_prob(tf.zeros([3])))
def test_bijector_scalar_underlying_ildj(self): d = tfd.Normal(0., 1.) # Uses Identity bijector, ildj=0. bij = tfd.Sample(d, [1]).experimental_default_event_space_bijector() ildj = bij.inverse_log_det_jacobian(tf.zeros([1, 1]), event_ndims=1) self.assertAllEqual(0., ildj) ildj = bij.inverse_log_det_jacobian(tf.zeros([1, 1]), event_ndims=2) self.assertAllEqual(0., ildj)
def test_bijector_constant_underlying_ildj(self): d = tfb.Scale([2., 3.])(tfd.Normal([0., 0.], 1.)) bij = tfd.Sample(d, [3]).experimental_default_event_space_bijector() ildj = bij.inverse_log_det_jacobian(tf.zeros([2, 3]), event_ndims=1) self.assertAllClose(-np.log([2., 3.]) * 3, ildj) ildj = bij.inverse_log_det_jacobian(tf.zeros([2, 3]), event_ndims=2) self.assertAllClose(-np.log([2., 3.]).sum() * 3, ildj)
def test_mixed_scalar(self): s = tfd.Sample(tfd.Independent(tfd.Normal(loc=[0., 0], scale=1), 1), 3, validate_args=False) x = s.sample(4, seed=test_util.test_seed()) lp = s.log_prob(x) self.assertEqual((4, 3, 2), x.shape) self.assertEqual((4,), lp.shape)
def testDivergence(self): """Neals funnel with large step size.""" strm = tfp_test_util.test_seed_stream() neals_funnel = tfd.JointDistributionSequential( [ tfd.Normal(loc=0., scale=3.), # b0 lambda y: tfd.Sample( # pylint: disable=g-long-lambda tfd.Normal(loc=0., scale=tf.math.exp(y / 2)), sample_shape=9), ], validate_args=True ) @tf.function(autograph=False) def run_chain_and_get_divergence(): nchains = 5 init_states = neals_funnel.sample(nchains, seed=strm()) _, has_divergence = tfp.mcmc.sample_chain( num_results=100, kernel=tfp.mcmc.NoUTurnSampler( target_log_prob_fn=lambda *args: neals_funnel.log_prob(args), step_size=[1., 1.], parallel_iterations=1, seed=strm()), current_state=init_states, trace_fn=lambda _, pkr: pkr.has_divergence, parallel_iterations=1) return tf.reduce_sum(tf.cast(has_divergence, dtype=tf.int32)) divergence_count = self.evaluate(run_chain_and_get_divergence()) # Test that we observe a fair among of divergence. self.assertAllGreater(divergence_count, 100)
def test_namedtuple_sample_log_prob(self): Model = collections.namedtuple('Model', ['e', 'scale', 'loc', 'm', 'x']) # pylint: disable=invalid-name # pylint: disable=bad-whitespace model = Model( e = tfd.Independent(tfd.Exponential(rate=[100, 120]), 1), scale=lambda e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]), loc = tfd.Normal(loc=0, scale=2.), m = tfd.Normal, x =lambda m: tfd.Sample(tfd.Bernoulli(logits=m), 12)) # pylint: enable=bad-whitespace d = tfd.JointDistributionNamed(model, validate_args=True) self.assertEqual( ( ('e', ()), ('scale', ('e',)), ('loc', ()), ('m', ('loc', 'scale')), ('x', ('m',)), ), d.resolve_graph()) xs = d.sample(seed=test_util.test_seed()) self.assertLen(xs, 5) # We'll verify the shapes work as intended when we plumb these back into the # respective log_probs. ds, _ = d.sample_distributions(value=xs, seed=test_util.test_seed()) self.assertLen(ds, 5) self.assertIsInstance(ds.e, tfd.Independent) self.assertIsInstance(ds.scale, tfd.Gamma) self.assertIsInstance(ds.loc, tfd.Normal) self.assertIsInstance(ds.m, tfd.Normal) self.assertIsInstance(ds.x, tfd.Sample) # Static properties. self.assertAllEqual(Model(e=tf.float32, scale=tf.float32, loc=tf.float32, m=tf.float32, x=tf.int32), d.dtype) batch_shape_tensor_, event_shape_tensor_ = self.evaluate([ d.batch_shape_tensor(), d.event_shape_tensor()]) expected_batch_shape = Model(e=[], scale=[], loc=[], m=[], x=[]) for (expected, actual_tensorshape, actual_shape_tensor_) in zip( expected_batch_shape, d.batch_shape, batch_shape_tensor_): self.assertAllEqual(expected, actual_tensorshape) self.assertAllEqual(expected, actual_shape_tensor_) expected_event_shape = Model(e=[2], scale=[], loc=[], m=[], x=[12]) for (expected, actual_tensorshape, actual_shape_tensor_) in zip( expected_event_shape, d.event_shape, event_shape_tensor_): self.assertAllEqual(expected, actual_tensorshape) self.assertAllEqual(expected, actual_shape_tensor_) expected_jlp = sum(d.log_prob(x) for d, x in zip(ds, xs)) actual_jlp = d.log_prob(xs) self.assertAllClose(*self.evaluate([expected_jlp, actual_jlp]), atol=0., rtol=1e-4)
def test_doc_string_2(self): n = 2000 c = np.r_[np.zeros(n // 2), np.ones(n // 2)] mean_0, mean_1 = 0, 5 x = np.r_[np.random.randn(n // 2).astype(dtype=np.float32) + mean_0, np.random.randn(n // 2).astype(dtype=np.float32) + mean_1] shuffle_idxs = np.arange(n) np.random.shuffle(shuffle_idxs) x = x[shuffle_idxs] c = c[shuffle_idxs] seed = test_util.test_seed_stream() # Density estimation with MADE. made = tfb.AutoregressiveNetwork( params=2, hidden_units=[1], event_shape=(1, ), kernel_initializer=tfk.initializers.VarianceScaling(0.1, seed=seed() % 2**31), conditional=True, conditional_event_shape=(1, )) distribution = tfd.TransformedDistribution( distribution=tfd.Sample(tfd.Normal(loc=0., scale=1.), sample_shape=[1]), bijector=tfb.MaskedAutoregressiveFlow(made)) # Construct and fit model. x_ = tfkl.Input(shape=(1, ), dtype=tf.float32) c_ = tfkl.Input(shape=(1, ), dtype=tf.float32) log_prob_ = distribution.log_prob( x_, bijector_kwargs={"conditional_input": c_}) model = tfk.Model([x_, c_], log_prob_) model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.1), loss=lambda _, log_prob: -log_prob) batch_size = 25 model.fit(x=[x, c], y=np.zeros((n, 0), dtype=np.float32), batch_size=batch_size, epochs=3, steps_per_epoch=n // batch_size, shuffle=False, verbose=True) # Use the fitted distribution to sample condition on c = 1 n_samples = 1000 cond = 1 samples = distribution.sample((n_samples, ), bijector_kwargs={ "conditional_input": cond * np.ones((n_samples, 1)) }, seed=seed()) # Assert mean is close to conditional mean self.assertAllMeansClose(samples[..., 0], mean_1, axis=0, atol=1.)
def test_summary_statistic(self, attr): sample_shape = [5, 4] mvn = tfd.Independent(tfd.Normal(loc=tf.zeros([3, 2]), scale=1), 1) d = tfd.Sample(mvn, sample_shape, validate_args=True) self.assertEqual((3, ), d.batch_shape) expected_stat = (getattr(mvn, attr)()[:, tf.newaxis, tf.newaxis, :] * tf.ones([3, 5, 4, 2])) actual_stat = getattr(d, attr)() self.assertAllEqual(*self.evaluate([expected_stat, actual_stat]))
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 = tfd.Normal(0, .2).sample(seed=stream()) pois = tfd.Poisson(log_rate=log_rate) lp = maybe_jit( tfd.Sample(pois, n, experimental_use_kahan_sum=True).log_prob)(samps) pois64 = tfd.Poisson(log_rate=tf.cast(log_rate, tf.float64)) lp64 = tfd.Sample(pois64, n).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, 0-80% GPU --vary_seed runs w/o experimental_use_kahan_sum. self.assertAllClose(lp64, lp, rtol=0., atol=.01)
def test_bijector_shapes(self): d = tfd.Sample(tfd.Uniform(tf.zeros([5]), 1.), 2) b = d.experimental_default_event_space_bijector() self.assertEqual((2,), d.event_shape) self.assertEqual((2,), b.inverse_event_shape((2,))) self.assertEqual((2,), b.forward_event_shape((2,))) self.assertEqual((5, 2), b.forward_event_shape((5, 2))) self.assertEqual((5, 2), b.inverse_event_shape((5, 2))) self.assertEqual((3, 5, 2), b.inverse_event_shape((3, 5, 2))) self.assertEqual((3, 5, 2), b.forward_event_shape((3, 5, 2))) d = tfd.Sample(tfd.CholeskyLKJ(4, concentration=tf.ones([5])), 2) b = d.experimental_default_event_space_bijector() self.assertEqual((2, 4, 4), d.event_shape) dim = (4 * 3) // 2 self.assertEqual((5, 2, dim), b.inverse_event_shape((5, 2, 4, 4))) self.assertEqual((5, 2, 4, 4), b.forward_event_shape((5, 2, dim))) self.assertEqual((3, 5, 2, dim), b.inverse_event_shape((3, 5, 2, 4, 4))) self.assertEqual((3, 5, 2, 4, 4), b.forward_event_shape((3, 5, 2, dim)))
def get_base_distribution(flat_event_size, dtype=DEFAULT_FLOAT_DTYPE_TF): base_standard_dist = tfd.JointDistributionSequential([ tfd.Sample( tfd.Normal( loc=tf.constant(0.0, dtype=dtype), scale=tf.constant(1.0, dtype=dtype), ), s, ) for s in flat_event_size ]) return base_standard_dist
def test_everything_scalar(self): s = tfd.Sample(tfd.Normal(loc=0, scale=1), 5, validate_args=True) x = s.sample(seed=test_util.test_seed()) actual_lp = s.log_prob(x) # Sample.log_prob will reduce over event space, ie, dims [0, 2] # corresponding to sizes concat([[5], [2]]). expected_lp = tf.reduce_sum(s.distribution.log_prob(x), axis=0) x_, actual_lp_, expected_lp_ = self.evaluate([x, actual_lp, expected_lp]) self.assertEqual((5,), x_.shape) self.assertEqual((), actual_lp_.shape) self.assertAllClose(expected_lp_, actual_lp_, atol=0, rtol=1e-3)
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 testMVN(self, event_shape, shift, tril, dynamic_shape): if dynamic_shape and tf.executing_eagerly(): self.skipTest('Eager execution does not support dynamic shape.') as_tensor = tf.convert_to_tensor if dynamic_shape: as_tensor = lambda v, name: tf1.placeholder_with_default( # pylint: disable=g-long-lambda v, shape=None, name='dynamic_' + name) fake_mvn = tfd.TransformedDistribution( distribution=tfd.Sample( tfd.Normal(loc=as_tensor(0., name='loc'), scale=as_tensor(1., name='scale'), validate_args=True), sample_shape=as_tensor(np.int32(event_shape), name='event_shape')), bijector=tfb.Chain( [tfb.Shift(shift=as_tensor(shift, name='shift')), tfb.ScaleMatvecTriL(scale_tril=as_tensor(tril, name='scale_tril')) ]), validate_args=True) base_dist = fake_mvn.distribution expected_mean = tf.linalg.matvec( tril, tf.broadcast_to(base_dist.mean(), shift.shape)) + shift expected_cov = tf.linalg.matmul( tril, tf.matmul( tf.linalg.diag(tf.broadcast_to(base_dist.variance(), shift.shape)), tril, adjoint_b=True)) expected_batch_shape = ps.shape(expected_mean)[:-1] if dynamic_shape: self.assertAllEqual(tf.TensorShape(None), fake_mvn.event_shape) self.assertAllEqual(tf.TensorShape(None), fake_mvn.batch_shape) else: self.assertAllEqual(event_shape, fake_mvn.event_shape) self.assertAllEqual(expected_batch_shape, fake_mvn.batch_shape) # Ensure sample works by checking first, second moments. num_samples = 7e3 y = fake_mvn.sample(int(num_samples), seed=test_util.test_seed()) x = y[0:5, ...] self.assertAllClose(expected_mean, tf.reduce_mean(y, axis=0), atol=0.1, rtol=0.1) self.assertAllClose(expected_cov, tfp.stats.covariance(y, sample_axis=0), atol=0., rtol=0.1) self.assertAllEqual(event_shape, fake_mvn.event_shape_tensor()) self.assertAllEqual(expected_batch_shape, fake_mvn.batch_shape_tensor()) self.assertAllEqual( ps.concat([[5], expected_batch_shape, event_shape], axis=0), ps.shape(x)) self.assertAllClose(expected_mean, fake_mvn.mean())
def _make_kernel(**kwargs): running_variance = tfp.experimental.stats.RunningVariance.from_stats( num_samples=10., mean=tf.zeros(5), variance=tf.ones(5)) kernel = tfp.experimental.mcmc.PreconditionedHamiltonianMonteCarlo( target_log_prob_fn=tfd.Sample(tfd.Normal(0., 1.), 5).log_prob, num_leapfrog_steps=2, step_size=1.) kernel = tfp.experimental.mcmc.DiagonalMassMatrixAdaptation( inner_kernel=kernel, initial_running_variance=running_variance, **kwargs) return kernel
def test_transformed_exp(self): sample_shape = 3 mvn = tfd.Independent(tfd.Normal(loc=[0., 0], scale=1), 1) exp = tfb.Exp() def expected_lp(y): x = exp.inverse(y) # Ie, tf.random.normal([4, 3, 2]) fldj = exp.forward_log_det_jacobian(x, event_ndims=1) return tf.reduce_sum(mvn.log_prob(x) - fldj, axis=1) # Transform a Sample. d = tfd.TransformedDistribution(tfd.Sample(mvn, sample_shape, validate_args=True), bijector=exp) y = d.sample(4, seed=test_util.test_seed()) actual_lp = d.log_prob(y) self.assertAllEqual((4, ) + (sample_shape, ) + (2, ), y.shape) self.assertAllEqual((4, ), actual_lp.shape) # If `TransformedDistribution` didn't scale the jacobian by # `_sample_distribution_size`, then `scale_fldj` would need to be `False`. self.assertAllClose(*self.evaluate([expected_lp(y), actual_lp]), atol=0., rtol=1e-3) # Sample a Transform. d = tfd.Sample(tfd.TransformedDistribution(mvn, bijector=exp), sample_shape, validate_args=True) y = d.sample(4, seed=test_util.test_seed()) actual_lp = d.log_prob(y) self.assertAllEqual((4, ) + (sample_shape, ) + (2, ), y.shape) self.assertAllEqual((4, ), actual_lp.shape) # Regardless of whether `TransformedDistribution` scales the jacobian by # `_sample_distribution_size`, `scale_fldj` is `True`. self.assertAllClose(*self.evaluate([expected_lp(y), actual_lp]), atol=0., rtol=1e-3)
def test_variable_sample_shape_exception(self): loc = tf.Variable(tf.zeros([4, 5, 3]), shape=tf.TensorShape(None)) scale = tf.Variable(tf.ones([]), shape=tf.TensorShape(None)) sample_shape = tf.Variable([[1, 2]], shape=tf.TensorShape(None)) with self.assertRaisesWithPredicateMatch( Exception, 'Argument `sample_shape` must be either a scalar or a vector.'): dist = tfd.Sample( tfd.Independent(tfd.Logistic(loc=loc, scale=scale), reinterpreted_batch_ndims=1), sample_shape=sample_shape, validate_args=True) self.evaluate([v.initializer for v in dist.trainable_variables]) self.evaluate(dist.mean())
def test_variable_shape_change(self): loc = tf.Variable(tf.zeros([4, 5, 3]), shape=tf.TensorShape(None)) scale = tf.Variable(tf.ones([]), shape=tf.TensorShape(None)) # In real life, you'd really always want `sample_shape` to be # `trainable=False`. sample_shape = tf.Variable([1, 2], shape=tf.TensorShape(None)) dist = tfd.Sample( tfd.Independent(tfd.Logistic(loc=loc, scale=scale), reinterpreted_batch_ndims=1), sample_shape=sample_shape, validate_args=True) self.evaluate([v.initializer for v in dist.trainable_variables]) x = dist.mean() y = dist.sample([7, 2], seed=test_util.test_seed()) loss_x = -dist.log_prob(x) loss_0 = -dist.log_prob(0.) batch_shape = dist.batch_shape_tensor() event_shape = dist.event_shape_tensor() [x_, y_, loss_x_, loss_0_, batch_shape_, event_shape_] = self.evaluate([ x, y, loss_x, loss_0, batch_shape, event_shape]) self.assertAllEqual([4, 5, 1, 2, 3], x_.shape) self.assertAllEqual([7, 2, 4, 5, 1, 2, 3], y_.shape) self.assertAllEqual([4, 5], loss_x_.shape) self.assertAllEqual([4, 5], loss_0_.shape) self.assertAllEqual([4, 5], batch_shape_) self.assertAllEqual([1, 2, 3], event_shape_) self.assertLen(dist.trainable_variables, 3) with tf.control_dependencies([ loc.assign(tf.zeros([])), scale.assign(tf.ones([3, 1, 2])), sample_shape.assign(6), ]): x = dist.mean() y = dist.sample([7, 2], seed=test_util.test_seed()) loss_x = -dist.log_prob(x) loss_0 = -dist.log_prob(0.) batch_shape = dist.batch_shape_tensor() event_shape = dist.event_shape_tensor() [x_, y_, loss_x_, loss_0_, batch_shape_, event_shape_] = self.evaluate([ x, y, loss_x, loss_0, batch_shape, event_shape]) self.assertAllEqual([3, 1, 6, 2], x_.shape) self.assertAllEqual([7, 2, 3, 1, 6, 2], y_.shape) self.assertAllEqual([3, 1], loss_x_.shape) self.assertAllEqual([3, 1], loss_0_.shape) self.assertAllEqual([3, 1], batch_shape_) self.assertAllEqual([6, 2], event_shape_) self.assertLen(dist.trainable_variables, 3)
def one_term(event_shape, event_shape_tensor, batch_shape, batch_shape_tensor, dtype): if not tensorshape_util.is_fully_defined(event_shape): event_shape = event_shape_tensor result = tfd.Sample(tfd.Uniform(low=tf.constant(-2., dtype=dtype), high=tf.constant(2., dtype=dtype)), sample_shape=event_shape) if not tensorshape_util.is_fully_defined(batch_shape): batch_shape = batch_shape_tensor needs_bcast = True else: # Only batch broadcast when batch ndims > 0. needs_bcast = bool(tensorshape_util.as_list(batch_shape)) if needs_bcast: result = tfd.BatchBroadcast(result, batch_shape) return result
def test_everything_nonscalar(self): s = tfd.Sample( tfd.Independent(tfd.Normal(loc=tf.zeros([3, 2]), scale=1), 1), [5, 4], validate_args=True) x = s.sample([6, 1], seed=test_util.test_seed()) actual_lp = s.log_prob(x) # Sample.log_prob will reduce over event space, ie, dims [2, 3, 5] # corresponding to sizes concat([[5, 4], [2]]). expected_lp = tf.reduce_sum( s.distribution.log_prob(tf.transpose(a=x, perm=[0, 1, 3, 4, 2, 5])), axis=[2, 3]) x_, actual_lp_, expected_lp_ = self.evaluate([x, actual_lp, expected_lp]) self.assertEqual((6, 1, 3, 5, 4, 2), x_.shape) self.assertEqual((6, 1, 3), actual_lp_.shape) self.assertAllClose(expected_lp_, actual_lp_, atol=0, rtol=1e-3)