def testDistribution(self, dist_name, data):
dist = data.draw(dhps.distributions(dist_name=dist_name, enable_vars=False,
validate_args=False))
seed = test_util.test_seed(sampler_type='stateless')
sample_shape = [2, 1]
with tfp_hps.no_tf_rank_errors(), kernel_hps.no_pd_errors():
s1, lp1 = dist.experimental_sample_and_log_prob(sample_shape, seed=seed)
s2 = dist.sample(sample_shape, seed=seed)
self.assertAllClose(s1, s2, atol=1e-4)
# Sanity-check the log prob. The actual values may differ arbitrarily (if
# the `sample_and_log_prob` implementation is more stable) or be NaN, but
# they should at least have the same shape.
lp2 = dist.log_prob(s1)
self.assertAllEqual(lp1.shape, lp2.shape)
def testDistribution(self, dist_name, data):
if dist_name in WORKING_PRECISION_TEST_BLOCK_LIST:
self.skipTest('{} is blocked'.format(dist_name))
def eligibility_filter(name):
return name not in WORKING_PRECISION_TEST_BLOCK_LIST
dist = data.draw(dhps.distributions(
dist_name=dist_name, eligibility_filter=eligibility_filter,
enable_vars=False, validate_args=False))
hp.note('Trying distribution {}'.format(
self.evaluate_dict(dist.parameters)))
seed = test_util.test_seed()
with tfp_hps.no_tf_rank_errors():
samples = dist.sample(5, seed=seed)
self.assertIn(samples.dtype, [tf.float32, tf.int32])
self.assertEqual(dist.log_prob(samples).dtype, tf.float32)
def log_prob_function(dist, x):
return dist.log_prob(x)
dist64 = tf.nest.map_structure(
tensor_to_f64, tfe.as_composite(dist), expand_composites=True)
with tfp_hps.no_tf_rank_errors():
result64 = log_prob_function(dist64, tensor_to_f64(samples))
self.assertEqual(result64.dtype, tf.float64)
def testDistribution(self, dist_name, data):
if dist_name in NO_NANS_IN_SAMPLE_TEST_BLOCK_LIST:
self.skipTest('{} is blocked'.format(dist_name))
def eligibility_filter(name):
return name not in NO_NANS_IN_SAMPLE_TEST_BLOCK_LIST
dist = data.draw(
dhps.distributions(dist_name=dist_name,
enable_vars=False,
eligibility_filter=eligibility_filter))
hp.note('Trying distribution {}'.format(
self.evaluate_dict(dist.parameters)))
seed = test_util.test_seed(sampler_type='stateless')
with tfp_hps.no_tf_rank_errors():
s1 = self.evaluate(dist.sample(20, seed=seed))
self.assertAllEqual(np.zeros_like(s1), np.isnan(s1))
def testKernelGradient(self, kernel_name, data):
event_dim = data.draw(hps.integers(min_value=2, max_value=4))
feature_ndims = data.draw(hps.integers(min_value=1, max_value=2))
feature_dim = data.draw(hps.integers(min_value=2, max_value=4))
kernel, kernel_parameter_variable_names = data.draw(
kernel_hps.kernels(
kernel_name=kernel_name,
event_dim=event_dim,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
enable_vars=True))
# Check that variable parameters get passed to the kernel.variables
kernel_variables_names = [
v.name.strip('_0123456789:') for v in kernel.variables]
self.assertEqual(
set(kernel_parameter_variable_names),
set(kernel_variables_names))
example_ndims = data.draw(hps.integers(min_value=1, max_value=2))
input_batch_shape = data.draw(tfp_hps.broadcast_compatible_shape(
kernel.batch_shape))
xs = tf.identity(data.draw(kernel_hps.kernel_input(
batch_shape=input_batch_shape,
example_ndims=example_ndims,
feature_dim=feature_dim,
feature_ndims=feature_ndims)))
# Check that we pick up all relevant kernel parameters.
wrt_vars = [xs] + list(kernel.variables)
self.evaluate([v.initializer for v in kernel.variables])
with tf.GradientTape() as tape:
with tfp_hps.assert_no_excessive_var_usage(
'method `apply` of {}'.format(kernel)):
tape.watch(wrt_vars)
with tfp_hps.no_tf_rank_errors():
diag = kernel.apply(xs, xs, example_ndims=example_ndims)
grads = tape.gradient(diag, wrt_vars)
assert_no_none_grad(kernel, 'apply', wrt_vars, grads)
self.assertAllClose(
diag,
type(kernel)(**kernel._parameters).apply(
xs, xs, example_ndims=example_ndims))
def check_statistic(self, dist, statistic, expected_static_shape,
expected_dynamic_shape):
try:
with tfp_hps.no_tf_rank_errors():
result = getattr(dist, statistic)()
msg = 'Shape {} not compatible with expected {}.'.format(
result.shape, expected_static_shape)
self.assertTrue(
expected_static_shape.is_compatible_with(
tf.broadcast_static_shape(result.shape,
expected_static_shape)), msg)
self.assertAllEqual(
self.evaluate(expected_dynamic_shape),
self.evaluate(
tf.broadcast_dynamic_shape(tf.shape(result),
expected_dynamic_shape)))
except NotImplementedError:
pass
def testDistribution(self, dist_name, data):
if dist_name in NO_NANS_TEST_BLOCK_LIST:
self.skipTest('{} is blocked'.format(dist_name))
def eligibility_filter(name):
return name not in NO_NANS_TEST_BLOCK_LIST
dist = data.draw(
dhps.distributions(dist_name=dist_name,
enable_vars=False,
eligibility_filter=eligibility_filter))
samples = self.check_samples_not_nan(dist)
self.assume_loc_scale_ok(dist)
hp.note('Testing on samples {}'.format(samples))
with tfp_hps.no_tf_rank_errors():
lp = self.evaluate(dist.log_prob(samples))
self.assertAllEqual(np.zeros_like(lp), np.isnan(lp))
def _test_sample_and_log_prob(self, dist_name, dist):
seed = test_util.test_seed(sampler_type='stateless')
num_samples = 3
sample = self.evaluate(
tf.function(dist.sample, jit_compile=True)(num_samples, seed=seed))
hp.note('Trying distribution {}'.format(
self.evaluate_dict(dist.parameters)))
hp.note('Drew samples {}'.format(sample))
xla_lp = self.evaluate(
tf.function(dist.log_prob,
jit_compile=True)(tf.convert_to_tensor(sample)))
with tfp_hps.no_tf_rank_errors():
graph_lp = self.evaluate(dist.log_prob(sample))
self.assertAllClose(xla_lp,
graph_lp,
atol=XLA_LOGPROB_ATOL[dist_name],
rtol=XLA_LOGPROB_RTOL[dist_name])
def _test_slicing(self, data, kernel_name, kernel, feature_dim,
feature_ndims):
example_ndims = data.draw(hps.integers(min_value=0, max_value=2))
batch_shape = kernel.batch_shape
slices = data.draw(tfp_hps.valid_slices(batch_shape))
slice_str = 'kernel[{}]'.format(', '.join(
tfp_hps.stringify_slices(slices)))
# Make sure the slice string appears in Hypothesis' attempted example log
hp.note('Using slice ' + slice_str)
if not slices: # Nothing further to check.
return
sliced_zeros = np.zeros(batch_shape)[slices]
sliced_kernel = kernel[slices]
hp.note('Using sliced kernel {}.'.format(sliced_kernel))
hp.note('Using sliced zeros {}.'.format(sliced_zeros.shape))
# Check that slicing modifies batch shape as expected.
self.assertAllEqual(sliced_zeros.shape, sliced_kernel.batch_shape)
xs = tf.identity(
data.draw(
kernel_hps.kernel_input(batch_shape=[],
example_ndims=example_ndims,
feature_dim=feature_dim,
feature_ndims=feature_ndims)))
# Check that apply of sliced kernels executes.
with tfp_hps.no_tf_rank_errors():
results = self.evaluate(
kernel.apply(xs, xs, example_ndims=example_ndims))
hp.note('Using results shape {}.'.format(results.shape))
sliced_results = self.evaluate(
sliced_kernel.apply(xs, xs, example_ndims=example_ndims))
# Come up with the slices for apply (which must also include example dims).
apply_slices = (tuple(slices) if isinstance(
slices, collections.Sequence) else (slices, ))
apply_slices += tuple([slice(None)] * example_ndims)
# Check that sampling a sliced kernel produces the same shape as
# slicing the samples from the original.
self.assertAllClose(results[apply_slices], sliced_results)
def check_event_space_bijector_constrains(self, dist, data):
event_space_bijector = dist.experimental_default_event_space_bijector()
if event_space_bijector is None:
return
# Draw a sample shape
sample_shape = data.draw(tfp_hps.shapes())
inv_event_shape = event_space_bijector.inverse_event_shape(
tensorshape_util.concatenate(dist.batch_shape, dist.event_shape))
# Draw a shape that broadcasts with `[batch_shape, inverse_event_shape]`
# where `inverse_event_shape` is the event shape in the bijector's
# domain. This is the shape of `y` in R**n, such that
# x = event_space_bijector(y) has the event shape of the distribution.
# TODO(b/174778703): Actually draw broadcast compatible shapes.
batch_inv_event_compat_shape = inv_event_shape
# batch_inv_event_compat_shape = data.draw(
# tfp_hps.broadcast_compatible_shape(inv_event_shape))
# batch_inv_event_compat_shape = tensorshape_util.concatenate(
# (1,) * (len(inv_event_shape) - len(batch_inv_event_compat_shape)),
# batch_inv_event_compat_shape)
total_sample_shape = tensorshape_util.concatenate(
sample_shape, batch_inv_event_compat_shape)
# full_sample_batch_event_shape = tensorshape_util.concatenate(
# sample_shape, inv_event_shape)
y = data.draw(
tfp_hps.constrained_tensors(tfp_hps.identity_fn,
total_sample_shape.as_list()))
hp.note('Trying to constrain inputs {}'.format(y))
with tfp_hps.no_tf_rank_errors():
x = event_space_bijector(y)
hp.note('Got constrained samples {}'.format(x))
with tf.control_dependencies(dist._sample_control_dependencies(x)):
self.evaluate(tensor_util.identity_as_tensor(x))
def check_event_space_bijector_constrains(self, dist, data):
event_space_bijector = dist.experimental_default_event_space_bijector()
if event_space_bijector is None:
return
total_sample_shape = tensorshape_util.concatenate(
# Draw a sample shape
data.draw(tfp_hps.shapes()),
# Draw a shape that broadcasts with `[batch_shape, inverse_event_shape]`
# where `inverse_event_shape` is the event shape in the bijector's
# domain. This is the shape of `y` in R**n, such that
# x = event_space_bijector(y) has the event shape of the distribution.
data.draw(tfp_hps.broadcasting_shapes(
event_space_bijector.inverse_event_shape(
tensorshape_util.concatenate(
dist.batch_shape, dist.event_shape)), n=1))[0])
y = data.draw(
tfp_hps.constrained_tensors(
tfp_hps.identity_fn, total_sample_shape.as_list()))
with tfp_hps.no_tf_rank_errors():
x = event_space_bijector(y)
with tf.control_dependencies(dist._sample_control_dependencies(x)):
self.evaluate(tf.identity(x))
def _test_slicing(self, data, dist_name, dist):
strm = test_util.test_seed_stream()
batch_shape = dist.batch_shape
slices = data.draw(dhps.valid_slices(batch_shape))
slice_str = 'dist[{}]'.format(', '.join(dhps.stringify_slices(slices)))
# Make sure the slice string appears in Hypothesis' attempted example log
hp.note('Using slice ' + slice_str)
if not slices: # Nothing further to check.
return
sliced_zeros = np.zeros(batch_shape)[slices]
sliced_dist = dist[slices]
hp.note('Using sliced distribution {}.'.format(sliced_dist))
# Check that slicing modifies batch shape as expected.
self.assertAllEqual(sliced_zeros.shape, sliced_dist.batch_shape)
if not sliced_zeros.size:
# TODO(b/128924708): Fix distributions that fail on degenerate empty
# shapes, e.g. Multinomial, DirichletMultinomial, ...
return
# Check that sampling of sliced distributions executes.
with tfp_hps.no_tf_rank_errors():
samples = self.evaluate(dist.sample(seed=strm()))
sliced_dist_samples = self.evaluate(
sliced_dist.sample(seed=strm()))
# Come up with the slices for samples (which must also include event dims).
sample_slices = (tuple(slices) if isinstance(
slices, collections.Sequence) else (slices, ))
if Ellipsis not in sample_slices:
sample_slices += (Ellipsis, )
sample_slices += tuple([slice(None)] *
tensorshape_util.rank(dist.event_shape))
sliced_samples = samples[sample_slices]
# Report sub-sliced samples (on which we compare log_prob) to hypothesis.
hp.note('Sample(s) for testing log_prob ' + str(sliced_samples))
# Check that sampling a sliced distribution produces the same shape as
# slicing the samples from the original.
self.assertAllEqual(sliced_samples.shape, sliced_dist_samples.shape)
# Check that a sliced distribution can compute the log_prob of its own
# samples (up to numerical validation errors).
with tfp_hps.no_tf_rank_errors():
try:
lp = self.evaluate(dist.log_prob(samples))
except tf.errors.InvalidArgumentError:
# TODO(b/129271256): d.log_prob(d.sample()) should not fail
# validate_args checks.
# We only tolerate this case for the non-sliced dist.
return
sliced_lp = self.evaluate(sliced_dist.log_prob(sliced_samples))
# Check that the sliced dist's log_prob agrees with slicing the original's
# log_prob.
# This `hp.assume` is suppressing array sizes that cause the sliced and
# non-sliced distribution to follow different Eigen code paths. Those
# different code paths lead to arbitrarily large variations in the results
# at parameter settings that Hypothesis is all too good at finding. Since
# the purpose of this test is just to check that we got slicing right, those
# discrepancies are a distraction.
# TODO(b/140229057): Remove this `hp.assume`, if and when Eigen's numerics
# become index-independent.
all_packetized = (_all_packetized(dist)
and _all_packetized(sliced_dist)
and _all_packetized(samples)
and _all_packetized(sliced_samples))
hp.note('Packetization check {}'.format(all_packetized))
all_non_packetized = (_all_non_packetized(dist)
and _all_non_packetized(sliced_dist)
and _all_non_packetized(samples)
and _all_non_packetized(sliced_samples))
hp.note('Non-packetization check {}'.format(all_non_packetized))
hp.assume(all_packetized or all_non_packetized)
self.assertAllClose(lp[slices],
sliced_lp,
atol=SLICING_LOGPROB_ATOL[dist_name],
rtol=SLICING_LOGPROB_RTOL[dist_name])
def testVmap(self, dist_name, data):
dist = data.draw(dhps.distributions(
dist_name=dist_name, enable_vars=False,
validate_args=False)) # TODO(b/142826246): Enable validate_args.
with tfp_hps.no_tf_rank_errors():
self._test_vectorization(dist_name, dist)
def testCompositeTensor(self, dist_name, data):
dist = data.draw(
dhps.distributions(
dist_name=dist_name, enable_vars=False, validate_args=False))
with tfp_hps.no_tf_rank_errors():
self._test_sample_and_log_prob(dist_name, dist)
