def test_doc_string_images_case_2(self): # Generate fake images. images = np.random.choice([0, 1], size=(100, 8, 8, 3)) n, width, height, channels = images.shape # Reshape images to achieve desired autoregressivity. reshaped_images = np.transpose(np.reshape( images, [n, width * height, channels]), axes=[0, 2, 1]) made = tfb.AutoregressiveNetwork(params=1, event_shape=[width * height], hidden_units=[20, 20], activation="relu") # Density estimation with MADE. # # NOTE: Parameterize an autoregressive distribution over an event_shape of # [channels, width * height], with univariate Bernoulli conditional # distributions. distribution = tfd.Autoregressive( lambda x: tfd.Independent( # pylint: disable=g-long-lambda tfd.Bernoulli(logits=tf.unstack(made(x), axis=-1)[0], dtype=tf.float32), reinterpreted_batch_ndims=2), sample0=tf.zeros([channels, width * height], dtype=tf.float32)) # Construct and fit model. x_ = tfkl.Input(shape=(channels, width * height), 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 = 10 model.fit( x=reshaped_images, 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((7, channels, width * height), distribution.sample(7).shape) self.assertAllEqual((n, ), distribution.log_prob(reshaped_images).shape)
def testVariableNumSteps(self): def fn(sample=0.): return tfd.Normal(loc=tf.zeros_like(sample), scale=1.) num_steps = tf.Variable(4, dtype=tf.int64) self.evaluate(num_steps.initializer) ar = tfd.Autoregressive(fn, num_steps=num_steps, validate_args=True) sample = ar.sample(seed=test_util.test_seed()) log_prob = ar.log_prob(sample) self.assertAllEqual([], sample.shape) self.assertAllEqual([], log_prob.shape) [sample, log_prob] = self.evaluate([sample, log_prob]) self.assertAllEqual([], sample.shape) self.assertAllEqual([], log_prob.shape)
def testSampleAndLogProbConsistency(self): batch_shape = np.int32([]) event_size = 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) self.run_test_sample_consistent_log_prob( self.evaluate, ar, num_samples=int(1e6), radius=1., center=0., rtol=0.01, seed=test_util.test_seed())
def testStatefulDistFn(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) def dist_fn(s): return StatefulNormal(loc=s) ar = tfd.Autoregressive(dist_fn, sample0=tfd.Normal(0., 1.).sample( 7, seed=test_util.test_seed()), num_steps=7) with warnings.catch_warnings(record=True) as triggered: self.evaluate(ar.sample(seed=test_util.test_seed())) self.assertTrue( any('Falling back to stateful sampling for `distribution_fn(sample0)`' in str(warning.message) for warning in triggered)) num_steps = tf.Variable(9) self.evaluate(num_steps.initializer) with warnings.catch_warnings(record=True) as triggered: self.evaluate( ar.copy(num_steps=num_steps).sample( seed=test_util.test_seed())) self.assertTrue( any('Falling back to stateful sampling for `distribution_fn(sample0)`' in str(warning.message) for warning in triggered))
def testVariableNumStepsAndEventShape(self): loc = tf.Variable(np.zeros((4, 2)), shape=tf.TensorShape(None)) event_ndims = tf.Variable(1) def fn(sample=None): if sample is not None: loc_param = tf.broadcast_to(loc, shape=tf.shape(sample)) else: loc_param = loc return tfd.Independent(tfd.Normal(loc=loc_param, scale=1.), reinterpreted_batch_ndims=event_ndims) num_steps = tf.Variable(7) self.evaluate([v.initializer for v in [loc, num_steps, event_ndims]]) ar = tfd.Autoregressive(fn, num_steps=num_steps, validate_args=True) sample = self.evaluate(ar.sample(3, seed=test_util.test_seed())) self.assertAllEqual([3, 4, 2], sample.shape) self.assertAllEqual([2], self.evaluate(ar.event_shape_tensor())) self.assertAllEqual([4], self.evaluate(ar.batch_shape_tensor()))
def testBatchAndEventShape(self): loc = tf.Variable(np.zeros((5, 1, 3)), shape=tf.TensorShape(None)) event_ndims = tf.Variable(2) def fn(sample): return tfd.Independent( tfd.Normal(loc=loc + 0. * sample, scale=1.), reinterpreted_batch_ndims=tf.convert_to_tensor(event_ndims)) self.evaluate([v.initializer for v in [loc, event_ndims]]) zero = tf.convert_to_tensor(0., dtype=loc.dtype) ar = tfd.Autoregressive(fn, num_steps=7, sample0=zero, validate_args=True) # NOTE: `ar.event_shape` and `ar.batch_shape` are not known statically, # even though the output of `ar.distribution_fn(...)` has statically-known # event shape and batch shape. self.assertEqual(ar.batch_shape, tf.TensorShape(None)) self.assertEqual(ar.event_shape, tf.TensorShape(None)) self.assertAllEqual([5], self.evaluate(ar.batch_shape_tensor())) self.assertAllEqual([1, 3], self.evaluate(ar.event_shape_tensor())) if tf.executing_eagerly(): self.assertEqual(tf.TensorShape([5]), ar.distribution_fn(zero).batch_shape) self.assertEqual(tf.TensorShape([1, 3]), ar.distribution_fn(zero).event_shape) with tf.control_dependencies( [loc.assign(np.zeros((4, 7))), event_ndims.assign(1)]): self.assertAllEqual([4], self.evaluate(ar.batch_shape_tensor())) self.assertAllEqual([7], self.evaluate(ar.event_shape_tensor())) if tf.executing_eagerly(): self.assertEqual(tf.TensorShape([4]), ar.distribution_fn(zero).batch_shape) self.assertEqual(tf.TensorShape([7]), ar.distribution_fn(zero).event_shape)
def testSampleIndependenceWithoutSeedBug(self): # Under eager, there was a bug where subsequent samples for position 0 would # be erroneously independent of the first sample when a seed was not # provided. # A pithy example: A simple 1-Markov autoregressive sequence for which the # first frame is either 0 or 1 with equal probability and all subsequent # frames copy the previous frame. Thus the overall sequence-level # distribution is 0000 with probability 0.5 and 1111 with probability 0.5. 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) ar = tfd.Autoregressive(distribution_fn, sample0=tf.constant([2, 2, 2, 2]), num_steps=4) samps = self.evaluate(ar.sample(10)) for s in samps: self.assertIn(np.mean(s), (0., 1.), msg=str(s))
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.Affine(scale_tril=self._random_scale_tril(event_size)) 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( distribution=tfd.Normal(loc=0., scale=1.), bijector=ar_flow, event_shape=[event_size], batch_shape=batch_shape, 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)