def testNanFromGradsDontPropagate(self):
"""Test that update with NaN gradients does not cause NaN in results."""
if tf1.control_flow_v2_enabled():
self.skipTest('b/138796859')
if tf.executing_eagerly(): return
def _nan_log_prob_with_nan_gradient(x):
return np.nan * tf.reduce_sum(x)
initial_x = tf.linspace(0.01, 5, 10)
hmc = tfp.mcmc.HamiltonianMonteCarlo(
target_log_prob_fn=_nan_log_prob_with_nan_gradient,
step_size=2.,
num_leapfrog_steps=5)
updated_x, kernel_results = hmc.one_step(
current_state=initial_x,
previous_kernel_results=hmc.bootstrap_results(initial_x),
seed=test_util.test_seed())
initial_x_, updated_x_, log_accept_ratio_ = self.evaluate(
[initial_x, updated_x, kernel_results.log_accept_ratio])
acceptance_probs = np.exp(np.minimum(log_accept_ratio_, 0.))
logging.vlog(1, 'initial_x = {}'.format(initial_x_))
logging.vlog(1, 'updated_x = {}'.format(updated_x_))
logging.vlog(1, 'log_accept_ratio = {}'.format(log_accept_ratio_))
self.assertAllEqual(initial_x_, updated_x_)
self.assertEqual(acceptance_probs, 0.)
self.assertAllEqual([True], [
g is None for g in tf.gradients(
ys=kernel_results.proposed_results.grads_target_log_prob,
xs=initial_x)
])
self.assertAllFinite(
self.evaluate(tf.gradients(ys=updated_x, xs=initial_x)[0]))
def testCholeskyUpdateXLA(self):
self.skip_if_no_xla()
if not (tf1.control_flow_v2_enabled() or self.use_static_shape):
self.skipTest('TF1 broken')
cholesky_update_fun = tf.function(tfp.math.cholesky_update,
jit_compile=True)
self._testCholeskyUpdate(cholesky_update_fun)
def testNotReparameterized(self):
if tf1.control_flow_v2_enabled():
self.skipTest("b/138796859")
total_count = tf.constant(5.0)
concentration = tf.constant([0.1, 0.1, 0.1])
_, [grad_total_count, grad_concentration] = tfp.math.value_and_gradient(
lambda n, c: tfd.DirichletMultinomial(n, c).sample(100),
[total_count, concentration])
self.assertIsNone(grad_total_count)
self.assertIsNone(grad_concentration)
def testNotReparameterized(self):
if tf1.control_flow_v2_enabled():
self.skipTest('b/138796859')
total_count = tf.constant(5.0)
probs = tf.constant([0.2, 0.6])
_, [grad_total_count, grad_probs] = tfp.math.value_and_gradient(
lambda n, p: tfd.Multinomial(total_count=n, probs=p).sample(100),
[total_count, probs])
self.assertIsNone(grad_total_count)
self.assertIsNone(grad_probs)
def testNotReparameterized(self):
if tf1.control_flow_v2_enabled():
self.skipTest('b/138796859')
total_count = tf.constant(5.0)
concentration = tf.constant([0.1, 0.1, 0.1])
_, [grad_total_count, grad_concentration] = tfp.math.value_and_gradient(
lambda n, c: tfd.DirichletMultinomial(n, c, validate_args=True).sample( # pylint: disable=g-long-lambda
100, seed=test_util.test_seed()), [total_count, concentration])
self.assertIsNone(grad_total_count)
self.assertIsNone(grad_concentration)
def _testWhileLoop(self):
"""Shows misuse of `dtc._TensorCoercible(distribution)` in `tf.while_loop`.
Since `dtc._TensorCoercible(distribution)` only creates the `Tensor` on an
as-needed basis, care must be taken that the Tensor is created outside the
body of a `tf.while_loop`, if the result is going to be used outside of the
`tf.while_loop`. Although this is the case for any use of `tf.while_loop` we
write this unit-test as a reminder of the behavior.
"""
mean_ = 0.5
stddev_ = 2.
num_iter_ = 4
x = dtc._TensorCoercible(tfd.Normal(mean_, 0.75),
tfd.Distribution.mean)
# Note: in graph mode we can make the assertion not raise
# if we make sure to create the Tensor outside the loop. Ie,
# tf.convert_to_tensor(x)
def _body(iter_, d):
y = dtc._TensorCoercible(tfd.Normal(0, stddev_),
tfd.Distribution.stddev)
return [iter_ + 1, d + y + x]
_, mean_plus_numiter_times_stddev = tf.while_loop(
cond=lambda iter_, *args: iter_ < num_iter_,
body=_body,
loop_vars=[0, mean_])
if not tf.executing_eagerly():
# In graph mode we cannot access the cached value of x outside the
# tf.while_loop. To make this exception not occur, simply call
# `tf.convert_to_tensor(x)` prior to the `tf.while_loop`. Doing so will
# cause the value of `x` to exist both outside and inside the
# `tf.while_loop`.
if tf1.control_flow_v2_enabled():
error_regex = r'Tensor.*must be from the same graph as Tensor.*'
else:
error_regex = 'Cannot use.*in a while loop'
with self.assertRaisesRegexp(ValueError, error_regex):
_ = x + tf.constant(3.)
return
# Things work in Eager mode because it has regular Python semantics,
# ie, lexical scope rules.
self.assertAllEqual([
mean_ + 3.,
mean_ + num_iter_ * (mean_ + stddev_),
],
self.evaluate([
x + tf.constant(3.),
mean_plus_numiter_times_stddev,
]))
def var_to_tensors(var):
if resource_variable_ops.is_resource_variable(var):
if tf.control_flow_v2_enabled() and hasattr(
layer_collection.graph, 'captures'):
# TODO(b/143690035): Note that the "captures" property relies on an
# API which might change.
captures = layer_collection.graph.captures
return [h for vh, h in captures if vh is var.handle]
else:
return [var.handle]
if utils.is_reference_variable(var):
return [tf_ops.internal_convert_to_tensor(var, as_ref=True)]
raise ValueError('%s is not a recognized variable type.' % str(var))
def testNotReparameterized(self):
if tf1.control_flow_v2_enabled():
self.skipTest('b/138796859')
total_count = tf.constant(5.0)
probs = tf.constant([0.4, 0.6])
_, [grad_total_count, grad_probs] = tfp.math.value_and_gradient(
lambda n, p: tfd.Multinomial( # pylint: disable=g-long-lambda
total_count=n,
probs=p,
validate_args=True).sample(100, seed=test_util.test_seed()),
[total_count, probs])
self.assertIsNone(grad_total_count)
self.assertIsNone(grad_probs)
def testNotReparameterized(self):
if tf1.control_flow_v2_enabled():
self.skipTest('b/138796859')
n = tf.constant(5.0)
c1 = tf.constant([0.1, 0.1, 0.1])
c0 = tf.constant([0.3, 0.3, 0.3])
def f(n, c1, c0):
dist = tfd.BetaBinomial(n, c1, c0, validate_args=True)
return dist.sample(100, seed=test_util.test_seed())
_, [grad_n, grad_c1, grad_c0] = tfp.math.value_and_gradient(f, [n, c1, c0])
self.assertIsNone(grad_n)
self.assertIsNone(grad_c1)
self.assertIsNone(grad_c0)
def variable_uses(self, var):
"""Computes number of times a variable is used.
Args:
var: Variable or ResourceVariable instance.
Returns:
Number of times a variable is used within this subgraph.
Raises:
ValueError: If 'var' is not a variable type.
"""
def _add_tensor_consumers_to_set(tensor, consumers_set):
"""Finds consumers of a tensor and add them to the current consumers set.
"""
for consumer in set(tensor.consumers()):
# These are the type of ops which relay a tensor to other ops without
# doing anything to the tensor value, so recursively find the actual
# consumers.
if consumer.type in [
"Identity", "ReadVariableOp", "Enter", "ResourceGather"]:
for output in consumer.outputs:
_add_tensor_consumers_to_set(output, consumers_set)
else:
consumers_set.add(consumer)
consumers = set()
if resource_variable_ops.is_resource_variable(var):
if tf.control_flow_v2_enabled() and hasattr(self._graph, "captures"):
# TODO(b/143690035): Note that the "captures" property relies on an API
# which might change.
captures = self._graph.captures
for handle in [h for vh, h in captures if vh is var.handle]:
_add_tensor_consumers_to_set(handle, consumers)
else:
_add_tensor_consumers_to_set(var.handle, consumers)
elif is_reference_variable(var):
_add_tensor_consumers_to_set(var.value(), consumers)
else:
raise ValueError("%s does not appear to be a variable." % str(var))
return len(self._members.intersection(consumers))
def test_forecast_from_hmc(self):
if not tf1.control_flow_v2_enabled():
self.skipTest(
'test_forecast_from_hmc does not currently work with TF1')
# test that we can directly plug in the output of an HMC chain as
# the input to `forecast`, as done in the example, with no `sess.run` call.
num_results = 5
num_timesteps = 4
num_steps_forecast = 3
batch_shape = [1, 2]
observed_time_series = self._build_tensor(
np.random.randn(*(batch_shape + [num_timesteps])))
model = self._build_model(observed_time_series)
samples, _ = tfp.sts.fit_with_hmc(model,
observed_time_series,
num_results=num_results,
num_warmup_steps=2,
num_variational_steps=2)
forecast_dist = tfp.sts.forecast(model,
observed_time_series,
parameter_samples=samples,
num_steps_forecast=num_steps_forecast)
forecast_mean = forecast_dist.mean()[..., 0]
forecast_scale = forecast_dist.stddev()[..., 0]
sample_shape = [10]
forecast_samples = forecast_dist.sample(
sample_shape, seed=test_util.test_seed())[..., 0]
self.evaluate(tf1.global_variables_initializer())
forecast_mean_, forecast_scale_, forecast_samples_ = self.evaluate(
(forecast_mean, forecast_scale, forecast_samples))
self.assertAllEqual(forecast_mean_.shape,
batch_shape + [num_steps_forecast])
self.assertAllEqual(forecast_scale_.shape,
batch_shape + [num_steps_forecast])
self.assertAllEqual(forecast_samples_.shape,
sample_shape + batch_shape + [num_steps_forecast])