def assert_univariate_target_conservation(test, target_d, step_size):
# Sample count limited partly by memory reliably available on Forge. The test
# remains reasonable even if the nuts recursion limit is severely curtailed
# (e.g., 3 or 4 levels), so use that to recover some memory footprint and bump
# the sample count.
num_samples = int(5e4)
num_steps = 1
strm = test_util.test_seed_stream()
# We wrap the initial values in `tf.identity` to avoid broken gradients
# resulting from a bijector cache hit, since bijectors of the same
# type/parameterization now share a cache.
# TODO(b/72831017): Fix broken gradients caused by bijector caching.
initialization = tf.identity(target_d.sample([num_samples], seed=strm()))
@tf.function(autograph=False)
def run_chain():
nuts = tfp.experimental.mcmc.PreconditionedNoUTurnSampler(
target_d.log_prob,
step_size=step_size,
max_tree_depth=3,
unrolled_leapfrog_steps=2)
result = tfp.mcmc.sample_chain(
num_results=num_steps,
num_burnin_steps=0,
current_state=initialization,
trace_fn=None,
kernel=nuts,
seed=strm())
return result
result = run_chain()
test.assertAllEqual([num_steps, num_samples], result.shape)
answer = result[0]
check_cdf_agrees = st.assert_true_cdf_equal_by_dkwm(
answer, target_d.cdf, false_fail_rate=1e-6)
check_enough_power = assert_util.assert_less(
st.min_discrepancy_of_true_cdfs_detectable_by_dkwm(
num_samples, false_fail_rate=1e-6, false_pass_rate=1e-6), 0.025)
movement = tf.abs(answer - initialization)
test.assertAllEqual([num_samples], movement.shape)
# This movement distance (1 * step_size) was selected by reducing until 100
# runs with independent seeds all passed.
check_movement = assert_util.assert_greater_equal(
tf.reduce_mean(movement), 1 * step_size)
return (check_cdf_agrees, check_enough_power, check_movement)
def _parameter_control_dependencies(self, is_init):
assertions = []
sample_shape = None # Memoize concretization.
# Check valid shape.
ndims_ = tensorshape_util.rank(self.sample_shape.shape)
if is_init != (ndims_ is None):
msg = 'Argument `sample_shape` must be either a scalar or a vector.'
if ndims_ is not None:
if ndims_ > 1:
raise ValueError(msg)
elif self.validate_args:
if sample_shape is None:
sample_shape = tf.convert_to_tensor(self.sample_shape)
assertions.append(
assert_util.assert_less(tf.rank(sample_shape),
2,
message=msg))
# Check valid dtype.
if is_init: # No xor check because `dtype` cannot change.
dtype_ = self.sample_shape.dtype
if dtype_ is None:
if sample_shape is None:
sample_shape = tf.convert_to_tensor(self.sample_shape)
dtype_ = sample_shape.dtype
if dtype_util.base_dtype(dtype_) not in {tf.int32, tf.int64}:
raise TypeError(
'Argument `sample_shape` must be integer type; '
'saw {}.'.format(dtype_util.name(dtype_)))
# Check valid "value".
if is_init != tensor_util.is_ref(self.sample_shape):
sample_shape_ = tf.get_static_value(self.sample_shape)
msg = 'Argument `sample_shape` must have non-negative values.'
if sample_shape_ is not None:
if np.any(np.array(sample_shape_) < 0):
raise ValueError('{} Saw: {}'.format(msg, sample_shape_))
elif self.validate_args:
if sample_shape is None:
sample_shape = tf.convert_to_tensor(self.sample_shape)
assertions.append(
assert_util.assert_greater(sample_shape, -1, message=msg))
return assertions
def _mean(self):
concentration = tf.convert_to_tensor(self.concentration)
mixing_concentration = tf.convert_to_tensor(self.mixing_concentration)
mixing_rate = tf.convert_to_tensor(self.mixing_rate)
mean = concentration * mixing_rate / (mixing_concentration - 1.)
if self.allow_nan_stats:
return tf.where(mixing_concentration > 1., mean,
dtype_util.as_numpy_dtype(self.dtype)(np.nan))
else:
with tf.control_dependencies([
assert_util.assert_less(
tf.ones([], self.dtype),
mixing_concentration,
message=
'mean undefined when `mixing_concentration` <= 1'),
]):
return tf.identity(mean)
def testSample(self):
a = tf.constant(1.0)
b = tf.constant(2.0)
n = 500000
d = tfd.SigmoidBeta(concentration0=a, concentration1=b, validate_args=True)
samples = d.sample(n, seed=test_util.test_seed())
sample_values = self.evaluate(samples)
self.assertEqual(samples.shape, (n,))
self.assertEqual(sample_values.shape, (n,))
self.assertTrue(self._kstest(a, b, sample_values))
check_cdf_agrees = st.assert_true_cdf_equal_by_dkwm(
samples, d.cdf, false_fail_rate=1e-6)
self.evaluate(check_cdf_agrees)
check_enough_power = assert_util.assert_less(
st.min_discrepancy_of_true_cdfs_detectable_by_dkwm(
n, false_fail_rate=1e-6, false_pass_rate=1e-6), 0.01)
self.evaluate(check_enough_power)
def _mode(self):
concentration = tf.convert_to_tensor(self.concentration)
k = tf.cast(tf.shape(concentration)[-1], self.dtype)
total_concentration = tf.reduce_sum(concentration, axis=-1)
mode = (concentration - 1.) / (total_concentration[..., tf.newaxis] - k)
if self.allow_nan_stats:
return tf.where(
tf.reduce_all(concentration > 1., axis=-1, keepdims=True),
mode,
dtype_util.as_numpy_dtype(self.dtype)(np.nan))
assertions = [
assert_util.assert_less(
tf.ones([], self.dtype),
concentration,
message='Mode undefined when any concentration <= 1')
]
with tf.control_dependencies(assertions):
return tf.identity(mode)
def _mean(self):
alpha0 = tf.convert_to_tensor(self.concentration1_numerator)
beta0 = tf.convert_to_tensor(self.concentration0_numerator)
alpha1 = tf.convert_to_tensor(self.concentration1_denominator)
beta1 = tf.convert_to_tensor(self.concentration0_denominator)
mean = alpha0 * (alpha1 + beta1 - 1.) / ((alpha0 + beta0) * (alpha1 - 1.))
if self.allow_nan_stats:
assertions = []
else:
assertions = [assert_util.assert_less(
tf.ones([], self.dtype), alpha1,
message='mean undefined when any denominator_concentration <= 1')]
with tf.control_dependencies(assertions):
return tf.where(
alpha1 > 1.,
mean,
dtype_util.as_numpy_dtype(self.dtype)(np.nan))
def _variance(self):
concentration = tf.convert_to_tensor(self.concentration)
scale = tf.convert_to_tensor(self.scale)
var = (tf.square(scale) / tf.square(concentration - 1.) /
(concentration - 2.))
if self.allow_nan_stats:
assertions = []
else:
assertions = [
assert_util.assert_less(
tf.constant(2., dtype=self.dtype),
concentration,
message='variance undefined when any concentration <= 2')
]
with tf.control_dependencies(assertions):
return tf.where(concentration > 2., var,
dtype_util.as_numpy_dtype(self.dtype)(np.nan))
def _parameter_control_dependencies(self, is_init):
"""Checks the validity of the concentration parameter."""
assertions = []
# In init, we can always build shape and dtype checks because
# we assume shape doesn't change for Variable backed args.
if is_init:
if not dtype_util.is_floating(self.concentration.dtype):
raise TypeError('Argument `concentration` must be float type.')
msg = 'Argument `concentration` must have rank at least 1.'
ndims = tensorshape_util.rank(self.concentration.shape)
if ndims is not None:
if ndims < 1:
raise ValueError(msg)
elif self.validate_args:
assertions.append(
assert_util.assert_rank_at_least(self.concentration,
1,
message=msg))
msg = 'Argument `concentration` must have `event_size` at least 2.'
event_size = tf.compat.dimension_value(
self.concentration.shape[-1])
if event_size is not None:
if event_size < 2:
raise ValueError(msg)
elif self.validate_args:
assertions.append(
assert_util.assert_less(1,
tf.shape(self.concentration)[-1],
message=msg))
if not self.validate_args:
assert not assertions # Should never happen.
return []
if is_init != tensor_util.is_ref(self.concentration):
assertions.append(
assert_util.assert_positive(
self.concentration,
message='Argument `concentration` must be positive.'))
return assertions
def _parameter_control_dependencies(self, is_init):
if not self.validate_args:
return []
if is_init:
try:
self._batch_shape()
except ValueError:
raise ValueError(
'Arguments `total_count`, `low` and `high` must have compatible '
'shapes; total_count.shape={}, low.shape={}, '
'high.shape={}.'.format(tf.shape(self.total_count),
tf.shape(self.low),
tf.shape(self.high)))
assertions = []
if is_init != tensor_util.is_ref(self.total_count):
total_count = tf.convert_to_tensor(self.total_count)
limit = BATES_TOTAL_COUNT_STABILITY_LIMITS[self.dtype]
msg = '`total_count` must be representable as a 32-bit integer.'
assertions.extend([
assert_util.assert_positive(
total_count, message='`total_count` must be positive.'),
distribution_util.assert_casting_closed(total_count,
target_dtype=tf.int32,
message=msg),
assert_util.assert_less_equal(
tf.cast(total_count, self.dtype),
tf.cast(limit, self.dtype),
message='`total_count` > {} is numerically unstable.'.
format(limit))
])
if is_init != (tensor_util.is_ref(self.low)
or tensor_util.is_ref(self.high)):
assertions.append(
assert_util.assert_less(
self.low,
self.high,
message='`low` must be less than `high`.'))
return assertions
def _mean(self):
df = tf.convert_to_tensor(self.df)
loc = tf.convert_to_tensor(self.loc)
scale = tf.convert_to_tensor(self.scale)
log_correction = (
tf.math.log(scale) + np.log(2.) + 0.5 *
(tf.math.log(df) - np.log(np.pi)) + tf.math.lgamma(0.5 * (df + 1.)) -
tf.math.lgamma(0.5 * df) - tf.math.log(df - 1))
mean = tf.math.exp(log_correction) + loc
if self.allow_nan_stats:
return tf.where(df > 1., mean,
dtype_util.as_numpy_dtype(self.dtype)(np.nan))
else:
return distribution_util.with_dependencies([
assert_util.assert_less(
tf.ones([], dtype=self.dtype),
df,
message='mean not defined for components of df <= 1'),
], mean)
def _variance(self):
concentration = tf.convert_to_tensor(self.concentration)
lim = tf.constant(.5, self.dtype)
valid = concentration < lim
safe_conc = tf.where(valid, concentration,
tf.constant(.25, self.dtype))
result = lambda: self.scale**2 / ((1 - safe_conc)**2 *
(1 - 2 * safe_conc))
if self.allow_nan_stats:
return tf.where(valid, result(),
tf.constant(float('nan'), self.dtype))
with tf.control_dependencies([
assert_util.assert_less(
concentration,
lim,
message=
'`variance` is undefined when `concentration >= 0.5`')
]):
return result()
def _mode(self):
concentration = tf.convert_to_tensor(self.concentration)
rate = tf.convert_to_tensor(self.rate)
mode = (concentration - 1.) / rate
if self.allow_nan_stats:
assertions = []
else:
assertions = [
assert_util.assert_less(
tf.ones([], self.dtype),
concentration,
message="Mode not defined when any concentration <= 1.")
]
with tf.control_dependencies(assertions):
nan = tf.fill(self._batch_shape_tensor(concentration=concentration,
rate=rate),
dtype_util.as_numpy_dtype(self.dtype)(np.nan),
name="nan")
return tf.where(concentration > 1., mode, nan)
def _variance(self):
concentration = tf.convert_to_tensor(self.concentration)
mixing_concentration = tf.convert_to_tensor(self.mixing_concentration)
mixing_rate = tf.convert_to_tensor(self.mixing_rate)
variance = (tf.square(concentration * mixing_rate /
(mixing_concentration - 1.)) /
(mixing_concentration - 2.))
if self.allow_nan_stats:
return tf.where(mixing_concentration > 2., variance,
dtype_util.as_numpy_dtype(self.dtype)(np.nan))
else:
with tf.control_dependencies([
assert_util.assert_less(
tf.ones([], self.dtype) * 2.,
mixing_concentration,
message=
'variance undefined when `mixing_concentration` <= 2')
]):
return tf.identity(variance)
def _parameter_control_dependencies(self, is_init):
if not self.validate_args:
return []
sample_shape = tf.concat(
[self._batch_shape_tensor(), self._event_shape_tensor()], axis=0)
low = None if self._low is None else tf.convert_to_tensor(self._low)
high = None if self._high is None else tf.convert_to_tensor(self._high)
assertions = []
if self._low is not None and is_init != tensor_util.is_ref(self._low):
low_shape = ps.shape(low)
broadcast_shape = ps.broadcast_shape(sample_shape, low_shape)
assertions.extend(
[distribution_util.assert_integer_form(
low, message='`low` has non-integer components.'),
assert_util.assert_equal(
tf.reduce_prod(broadcast_shape),
tf.reduce_prod(sample_shape),
message=('Shape of `low` adds extra batch dimensions to '
'sample shape.'))])
if self._high is not None and is_init != tensor_util.is_ref(self._high):
high_shape = ps.shape(high)
broadcast_shape = ps.broadcast_shape(sample_shape, high_shape)
assertions.extend(
[distribution_util.assert_integer_form(
high, message='`high` has non-integer components.'),
assert_util.assert_equal(
tf.reduce_prod(broadcast_shape),
tf.reduce_prod(sample_shape),
message=('Shape of `high` adds extra batch dimensions to '
'sample shape.'))])
if (self._low is not None and self._high is not None and
(is_init != (tensor_util.is_ref(self._low)
or tensor_util.is_ref(self._high)))):
assertions.append(assert_util.assert_less(
low, high,
message='`low` must be strictly less than `high`.'))
return assertions
def assert_univariate_target_conservation(test, target_d, step_size):
# Sample count limited partly by memory reliably available on Forge. The test
# remains reasonable even if the nuts recursion limit is severely curtailed
# (e.g., 3 or 4 levels), so use that to recover some memory footprint and bump
# the sample count.
num_samples = int(5e4)
num_steps = 1
strm = tfp.util.SeedStream(salt='univariate_nuts_test', seed=1)
initialization = target_d.sample([num_samples], seed=strm())
@tf.function(autograph=False)
def run_chain():
nuts = tfp.mcmc.NoUTurnSampler(
target_d.log_prob,
step_size=step_size,
max_tree_depth=3,
unrolled_leapfrog_steps=2,
seed=strm())
result, _ = tfp.mcmc.sample_chain(
num_results=num_steps,
num_burnin_steps=0,
current_state=initialization,
kernel=nuts)
return result
result = run_chain()
test.assertAllEqual([num_steps, num_samples], result.shape)
answer = result[0]
check_cdf_agrees = st.assert_true_cdf_equal_by_dkwm(
answer, target_d.cdf, false_fail_rate=1e-6)
check_enough_power = assert_util.assert_less(
st.min_discrepancy_of_true_cdfs_detectable_by_dkwm(
num_samples, false_fail_rate=1e-6, false_pass_rate=1e-6), 0.025)
movement = tf.abs(answer - initialization)
test.assertAllEqual([num_samples], movement.shape)
# This movement distance (1 * step_size) was selected by reducing until 100
# runs with independent seeds all passed.
check_movement = assert_util.assert_greater_equal(
tf.reduce_mean(movement), 1 * step_size)
return (check_cdf_agrees, check_enough_power, check_movement)
def _parameter_control_dependencies(self, is_init):
if not self.validate_args:
return []
assertions = []
tailweight_is_ref = tensor_util.is_ref(self.tailweight)
tailweight = tf.convert_to_tensor(self.tailweight)
if (is_init != tailweight_is_ref and
is_init != tensor_util.is_ref(self.skewness)):
assertions.append(assert_util.assert_less(
tf.math.abs(self.skewness),
tailweight,
message='Expect `tailweight > |skewness|`'))
if is_init != tensor_util.is_ref(self.scale):
assertions.append(assert_util.assert_positive(
self.scale,
message='Argument `scale` must be positive.'))
if is_init != tailweight_is_ref:
assertions.append(assert_util.assert_positive(
tailweight,
message='Argument `tailweight` must be positive.'))
return assertions
def testMean(self, dtype):
testee_lkj = tfd.LKJ(dimension=3, concentration=dtype([1., 3., 5.]))
num_samples = 20000
results = testee_lkj.sample(sample_shape=[num_samples])
mean = testee_lkj.mean()
self.assertEqual(mean.shape, [3, 3, 3])
check1 = st.assert_true_mean_equal_by_dkwm(
samples=results, low=-1., high=1.,
expected=mean,
false_fail_rate=1e-6)
check2 = assert_util.assert_less(
st.min_discrepancy_of_true_means_detectable_by_dkwm(
num_samples,
low=-1.,
high=1.,
# Smaller false fail rate because of different batch sizes between
# these two checks.
false_fail_rate=1e-7,
false_pass_rate=1e-6),
# 4% relative error
0.08)
self.evaluate([check1, check2])
def _parameter_control_dependencies(self, is_init):
if not self.validate_args:
return []
low = tf.convert_to_tensor(self.low)
high = tf.convert_to_tensor(self.high)
peak = tf.convert_to_tensor(self.peak)
assertions = []
if (is_init != tensor_util.is_ref(self.low) and
is_init != tensor_util.is_ref(self.high)):
assertions.append(assert_util.assert_less(
low, high, message='triangular not defined when low >= high.'))
if (is_init != tensor_util.is_ref(self.low) and
is_init != tensor_util.is_ref(self.peak)):
assertions.append(
assert_util.assert_less_equal(
low, peak, message='triangular not defined when low > peak.'))
if (is_init != tensor_util.is_ref(self.high) and
is_init != tensor_util.is_ref(self.peak)):
assertions.append(
assert_util.assert_less_equal(
peak, high, message='triangular not defined when peak > high.'))
return assertions
def _variance(self):
concentration = tf.convert_to_tensor(self.concentration)
lim = tf.constant(.5, self.dtype)
valid = concentration < lim
safe_conc = tf.where(valid, concentration,
tf.constant(.25, self.dtype))
def result():
answer = self.scale**2 / ((1 - safe_conc)**2 * (1 - 2 * safe_conc))
# Force broadcasting with self.loc to get the shape right, even though the
# variance doesn't depend on the location.
return answer + tf.zeros_like(self.loc)
if self.allow_nan_stats:
return tf.where(valid, result(),
tf.constant(float('nan'), self.dtype))
with tf.control_dependencies([
assert_util.assert_less(
concentration,
lim,
message=
'`variance` is undefined when `concentration >= 0.5`')
]):
return result()
def _variance(self):
df = tf.convert_to_tensor(self.df)
scale = tf.convert_to_tensor(self.scale)
# We need to put the tf.where inside the outer tf.where to ensure we never
# hit a NaN in the gradient.
denom = tf.where(df > 2., df - 2., tf.ones_like(df))
# Abs(scale) superfluous.
var = (tf.ones(self._batch_shape_tensor(df=df, scale=scale),
dtype=self.dtype) * tf.square(scale) * df / denom)
# When 1 < df <= 2, variance is infinite.
result_where_defined = tf.where(
df > 2., var,
dtype_util.as_numpy_dtype(self.dtype)(np.inf))
if self.allow_nan_stats:
return tf.where(df > 1., result_where_defined,
dtype_util.as_numpy_dtype(self.dtype)(np.nan))
else:
return distribution_util.with_dependencies([
assert_util.assert_less(
tf.ones([], dtype=self.dtype),
df,
message='variance not defined for components of df <= 1'),
], result_where_defined)
def kendalls_tau(y_true, y_pred, name=None):
"""Computes Kendall's Tau for two ordered lists.
Kendall's Tau measures the correlation between ordinal rankings. This
implementation is similar to the one used in scipy.stats.kendalltau.
The provided values may be of any type that is sortable, with the
argsort indices indicating the true or proposed ordinal sequence.
Args:
y_true: a `Tensor` of shape `[n]` containing the true ordinal ranking.
y_pred: a `Tensor` of shape `[n]` containing the predicted ordering of the
same N items.
name: Optional Python `str` name for ops created by this method.
Default value: `None` (i.e., 'kendalls_tau').
Returns:
kendalls_tau: Kendall's Tau, the 1945 tau-b formulation that ignores
ordering of ties, as a `float32` scalar Tensor.
"""
with tf.name_scope(name or 'kendalls_tau'):
in_type = dtype_util.common_dtype([y_true, y_pred],
dtype_hint=tf.float32)
y_true = tf.convert_to_tensor(y_true, name='y_true', dtype=in_type)
y_pred = tf.convert_to_tensor(y_pred, name='y_pred', dtype=in_type)
tensorshape_util.assert_is_compatible_with(y_true.shape, y_pred.shape)
assertions = [
assert_util.assert_rank(y_true, 1),
assert_util.assert_greater(
ps.size(y_true), 1, 'Ordering requires at least 2 elements.')
]
with tf.control_dependencies(assertions):
lexa = lexicographical_indirect_sort(y_true, y_pred)
# See A Computer Method for Calculating Kendall's Tau with Ungrouped Data
# by William Night, Journal of the American Statistical Association,
# Jun., 1966, Vol. 61, No. 314, Part 1 (Jun., 1966), pp. 436-439
# for notation https://www.jstor.org/stable/2282833
def jointly_tied_pairs_body(first, t, i):
not_equal = tf.math.logical_or(
tf.not_equal(y_true[lexa[first]], y_true[lexa[i]]),
tf.not_equal(y_pred[lexa[first]], y_pred[lexa[i]]))
return (tf.where(not_equal, i, first),
tf.where(not_equal,
t + ((i - first) * (i - first - 1)) // 2,
t), i + 1)
n = ps.size0(y_true)
first, t, _ = tf.while_loop(cond=lambda first, t, i: i < n,
body=jointly_tied_pairs_body,
loop_vars=(0, 0, 1))
t += ((n - first) * (n - first - 1)) // 2
def ties_y_true_body(first, v, i):
not_equal = tf.not_equal(y_true[lexa[first]], y_true[lexa[i]])
return (tf.where(not_equal, i, first),
tf.where(not_equal,
v + ((i - first) * (i - first - 1)) // 2,
v), i + 1)
first, v, _ = tf.while_loop(cond=lambda first, v, i: i < n,
body=ties_y_true_body,
loop_vars=(0, 0, 1))
v += ((n - first) * (n - first - 1)) // 2
# count exchanges
exchanges, newperm = iterative_mergesort(y_pred, lexa)
def ties_in_y_pred_body(first, u, i):
not_equal = tf.not_equal(y_pred[newperm[first]],
y_pred[newperm[i]])
return (tf.where(not_equal, i, first),
tf.where(not_equal,
u + ((i - first) * (i - first - 1)) // 2,
u), i + 1)
first, u, _ = tf.while_loop(cond=lambda first, u, i: i < n,
body=ties_in_y_pred_body,
loop_vars=(0, 0, 1))
u += ((n - first) * (n - first - 1)) // 2
n0 = (n * (n - 1)) // 2
assertions = [
assert_util.assert_less(v, tf.cast(n0, tf.int32),
'All ranks are ties for y_true.'),
assert_util.assert_less(u, tf.cast(n0, tf.int32),
'All ranks are ties for y_pred.')
]
with tf.control_dependencies(assertions):
return (tf.cast(n0 - (u + v - t), tf.float32) -
2.0 * tf.cast(exchanges, tf.float32)) / tf.math.sqrt(
tf.cast(n0 - v, tf.float32) *
tf.cast(n0 - u, tf.float32))
def __init__(self,
low=0.,
high=1.,
peak=0.5,
validate_args=False,
allow_nan_stats=True,
name="Triangular"):
"""Initialize a batch of Triangular distributions.
Args:
low: Floating point tensor, lower boundary of the output interval. Must
have `low < high`.
Default value: `0`.
high: Floating point tensor, upper boundary of the output interval. Must
have `low < high`.
Default value: `1`.
peak: Floating point tensor, mode of the output interval. Must have
`low <= peak` and `peak <= high`.
Default value: `0.5`.
validate_args: Python `bool`, default `False`. When `True` distribution
parameters are checked for validity despite possibly degrading runtime
performance. When `False` invalid inputs may silently render incorrect
outputs.
Default value: `False`.
allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
(e.g., mean, mode, variance) use the value "`NaN`" to indicate the
result is undefined. When `False`, an exception is raised if one or
more of the statistic's batch members are undefined.
Default value: `True`.
name: Python `str` name prefixed to Ops created by this class.
Default value: `'Triangular'`.
Raises:
InvalidArgumentError: if `validate_args=True` and one of the following is
True:
* `low >= high`.
* `peak > high`.
* `low > peak`.
"""
parameters = locals()
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([low, high, peak], tf.float32)
low = tf.convert_to_tensor(value=low, name="low", dtype=dtype)
high = tf.convert_to_tensor(value=high, name="high", dtype=dtype)
peak = tf.convert_to_tensor(value=peak, name="peak", dtype=dtype)
with tf.control_dependencies([
assert_util.assert_less(
low, high, message="triangular not defined when low >= high."),
assert_util.assert_less_equal(
low, peak, message="triangular not defined when low > peak."),
assert_util.assert_less_equal(
peak, high, message="triangular not defined when peak > high."),
] if validate_args else []):
self._low = tf.identity(low, name="low")
self._high = tf.identity(high, name="high")
self._peak = tf.identity(peak, name="peak")
tf.debugging.assert_same_float_dtype(
[self._low, self._high, self._peak])
super(Triangular, self).__init__(
dtype=self._low.dtype,
reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
graph_parents=[self._low, self._high, self._peak],
name=name)
def __init__(self,
distribution,
low=None,
high=None,
validate_args=False,
name="QuantizedDistribution"):
"""Construct a Quantized Distribution representing `Y = ceiling(X)`.
Some properties are inherited from the distribution defining `X`. Example:
`allow_nan_stats` is determined for this `QuantizedDistribution` by reading
the `distribution`.
Args:
distribution: The base distribution class to transform. Typically an
instance of `Distribution`.
low: `Tensor` with same `dtype` as this distribution and shape
able to be added to samples. Should be a whole number. Default `None`.
If provided, base distribution's `prob` should be defined at
`low`.
high: `Tensor` with same `dtype` as this distribution and shape
able to be added to samples. Should be a whole number. Default `None`.
If provided, base distribution's `prob` should be defined at
`high - 1`.
`high` must be strictly greater than `low`.
validate_args: Python `bool`, default `False`. When `True` distribution
parameters are checked for validity despite possibly degrading runtime
performance. When `False` invalid inputs may silently render incorrect
outputs.
name: Python `str` name prefixed to Ops created by this class.
Raises:
TypeError: If `dist_cls` is not a subclass of
`Distribution` or continuous.
NotImplementedError: If the base distribution does not implement `cdf`.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
self._dist = distribution
if low is not None:
low = tf.convert_to_tensor(low,
name="low",
dtype=distribution.dtype)
if high is not None:
high = tf.convert_to_tensor(high,
name="high",
dtype=distribution.dtype)
dtype_util.assert_same_float_dtype(
tensors=[self.distribution, low, high])
# We let QuantizedDistribution access _graph_parents since this class is
# more like a baseclass.
graph_parents = self._dist._graph_parents # pylint: disable=protected-access
checks = []
if validate_args and low is not None and high is not None:
message = "low must be strictly less than high."
checks.append(
assert_util.assert_less(low, high, message=message))
self._validate_args = validate_args # self._check_integer uses this.
with tf.control_dependencies(checks if validate_args else []):
if low is not None:
self._low = self._check_integer(low)
graph_parents += [self._low]
else:
self._low = None
if high is not None:
self._high = self._check_integer(high)
graph_parents += [self._high]
else:
self._high = None
super(QuantizedDistribution, self).__init__(
dtype=self._dist.dtype,
reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=self._dist.allow_nan_stats,
parameters=parameters,
graph_parents=graph_parents,
name=name)
def _sample_n(self, n, seed=None):
dim0_seed, otherdims_seed = samplers.split_seed(seed,
salt='von_mises_fisher')
# The sampling strategy relies on the fact that vMF variates are symmetric
# about the mean direction. Accordingly, if we have a sampling strategy for
# the away-from-mean angle, then we can uniformly sample the remaining
# dimensions on the S^{dim-2} sphere for , and rotate these samples from a
# (1, 0, 0, ..., 0)-mode distribution into the target orientation.
#
# This is easy to imagine on the 1-sphere (S^1; in 2-D space): sample a
# von-Mises distributed `x` value in [-1, 1], then uniformly select what
# amounts to a "up" or "down" additional degree of freedom after unit
# normalizing, followed by a final rotation to the desired mean direction
# from a basis of (1, 0).
#
# On S^2 (in 3-D), selecting a vMF `x` identifies a circle in `yz` on the
# unit sphere over which the distribution is uniform, in particular the
# circle where x = \hat{x} intersects the unit sphere. We pick a point on
# that circle, then rotate to the desired mean direction from a basis of
# (1, 0, 0).
mean_direction = tf.convert_to_tensor(self.mean_direction)
concentration = tf.convert_to_tensor(self.concentration)
event_dim = (
tf.compat.dimension_value(self.event_shape[0]) or
self._event_shape_tensor(mean_direction=mean_direction)[0])
sample_batch_shape = ps.concat([[n], self._batch_shape_tensor(
mean_direction=mean_direction, concentration=concentration)], axis=0)
dim = tf.cast(event_dim - 1, self.dtype)
if event_dim == 3:
samples_dim0 = self._sample_3d(n,
mean_direction=mean_direction,
concentration=concentration,
seed=dim0_seed)
else:
# Wood'94 provides a rejection algorithm to sample the x coordinate.
# Wood'94 definition of b:
# b = (-2 * kappa + tf.sqrt(4 * kappa**2 + dim**2)) / dim
# https://stats.stackexchange.com/questions/156729 suggests:
b = dim / (2 * concentration +
tf.sqrt(4 * concentration**2 + dim**2))
# TODO(bjp): Integrate any useful numerical tricks from hyperspherical VAE
# https://github.com/nicola-decao/s-vae-tf/
x = (1 - b) / (1 + b)
c = concentration * x + dim * tf.math.log1p(-x**2)
beta = beta_lib.Beta(dim / 2, dim / 2)
def cond_fn(w, should_continue, seed):
del w, seed
return tf.reduce_any(should_continue)
def body_fn(w, should_continue, seed):
"""While loop body for sampling the angle `w`."""
beta_seed, unif_seed, next_seed = samplers.split_seed(seed, n=3)
z = beta.sample(sample_shape=sample_batch_shape, seed=beta_seed)
# set_shape needed here because of b/139013403
tensorshape_util.set_shape(z, w.shape)
w = tf.where(should_continue,
(1. - (1. + b) * z) / (1. - (1. - b) * z),
w)
if not self.allow_nan_stats:
w = tf.debugging.check_numerics(w, 'w')
unif = samplers.uniform(
sample_batch_shape, seed=unif_seed, dtype=self.dtype)
# set_shape needed here because of b/139013403
tensorshape_util.set_shape(unif, w.shape)
should_continue = should_continue & (
concentration * w + dim * tf.math.log1p(-x * w) - c <
# Use log1p(-unif) to prevent log(0) and ensure that log(1) is
# possible.
tf.math.log1p(-unif))
return w, should_continue, next_seed
w = tf.zeros(sample_batch_shape, dtype=self.dtype)
should_continue = tf.ones(sample_batch_shape, dtype=tf.bool)
samples_dim0, _, _ = tf.while_loop(
cond=cond_fn, body=body_fn,
loop_vars=(w, should_continue, dim0_seed))
samples_dim0 = samples_dim0[..., tf.newaxis]
if not self._allow_nan_stats:
# Verify samples are w/in -1, 1, with useful error output tensors (top
# value rather than all values).
with tf.control_dependencies([
assert_util.assert_less_equal(
samples_dim0,
dtype_util.as_numpy_dtype(self.dtype)(1.01)),
assert_util.assert_greater_equal(
samples_dim0,
dtype_util.as_numpy_dtype(self.dtype)(-1.01)),
]):
samples_dim0 = tf.identity(samples_dim0)
samples_otherdims_shape = ps.concat([sample_batch_shape, [event_dim - 1]],
axis=0)
unit_otherdims = tf.math.l2_normalize(
samplers.normal(
samples_otherdims_shape, seed=otherdims_seed, dtype=self.dtype),
axis=-1)
samples = tf.concat([
samples_dim0, # we must avoid sqrt(1 - (>1)**2)
tf.sqrt(tf.maximum(1 - samples_dim0**2, 0.)) * unit_otherdims
], axis=-1)
samples = tf.math.l2_normalize(samples, axis=-1)
if not self.allow_nan_stats:
samples = tf.debugging.check_numerics(samples, 'samples')
# Runtime assert that samples are unit length.
if not self.allow_nan_stats:
worst, _ = tf.math.top_k(
tf.reshape(tf.abs(1 - tf.linalg.norm(samples, axis=-1)), [-1]))
with tf.control_dependencies([
assert_util.assert_near(
dtype_util.as_numpy_dtype(self.dtype)(0),
worst,
atol=1e-4,
summarize=100)
]):
samples = tf.identity(samples)
# The samples generated are symmetric around a mode at (1, 0, 0, ...., 0).
# Now, we move the mode to `self.mean_direction` using a rotation matrix.
if not self.allow_nan_stats:
# Assert that the basis vector rotates to the mean direction, as expected.
basis = tf.cast(tf.concat([[1.], tf.zeros([event_dim - 1])], axis=0),
self.dtype)
with tf.control_dependencies([
assert_util.assert_less(
tf.linalg.norm(
self._rotate(basis, mean_direction=mean_direction) -
mean_direction, axis=-1),
dtype_util.as_numpy_dtype(self.dtype)(1e-5))
]):
return self._rotate(samples, mean_direction=mean_direction)
return self._rotate(samples, mean_direction=mean_direction)
def _testSampleConsistentLogProbInterval(self,
concentrations,
det_bounds,
dim,
num_samples=int(1e5),
dtype=np.float32,
input_output_cholesky=False,
false_fail_rate=1e-6,
target_discrepancy=0.1,
seed=42):
# Consider the set M of dim x dim correlation matrices whose
# determinant exceeds some bound (rationale for bound forthwith).
# - This is a (convex!) shape in dim * (dim - 1) / 2 dimensions
# (because a correlation matrix is determined by its lower
# triangle, and the main diagonal is all 1s).
# - Further, M is contained entirely in the [-1,1] cube,
# because no correlation can fall outside that interval.
#
# We have two different ways to estimate the volume of M:
# - Importance sampling from the LKJ distribution
# - Importance sampling from the uniform distribution on the cube
#
# This test checks that these two methods agree. However, because
# the uniform proposal leads to many rejections (thus slowness),
# those volumes are computed offline and the confidence intervals
# are presented to this test procedure in the "volume_bounds"
# table.
#
# Why place a lower bound on the determinant? Because for eta > 1,
# the density of LKJ approaches 0 as the determinant approaches 0.
# However, the test methodology requires an upper bound on the
# improtance weights produced. Rejecting matrices with too-small
# determinant (from both methods) allows me to supply that bound.
#
# I considered several alternative regions whose volume I might
# know analytically (without having to do rejection).
# - Option a: Some hypersphere guaranteed to be contained inside M.
# - Con: I don't know a priori how to find a radius for it.
# - Con: I still need a lower bound on the determinants that appear
# in this sphere, and I don't know how to compute it.
# - Option b: Some trapezoid given as the convex hull of the
# nearly-extreme correlation matrices (i.e., those that partition
# the variables into two strongly anti-correclated groups).
# - Con: Would have to dig up n-d convex hull code to implement this.
# - Con: Need to compute the volume of that convex hull.
# - Con: Need a bound on the determinants of the matrices in that hull.
# - Option c: Same thing, but with the matrices that make a single pair
# of variables strongly correlated (or anti-correlated), and leaves
# the others uncorrelated.
# - Same cons, except that there is a determinant bound (which
# felt pretty loose).
lows = [dtype(volume_bounds[dim][db][0]) for db in det_bounds]
highs = [dtype(volume_bounds[dim][db][1]) for db in det_bounds]
concentration = np.array(concentrations, dtype=dtype)
det_bounds = np.array(det_bounds, dtype=dtype)
# Due to possible numerical inaccuracies while lower bounding the
# determinant, the maximum of the importance weights may exceed the
# theoretical maximum (importance_maxima). We add a tolerance to guard
# against this. An alternative would have been to add a threshold while
# filtering in _det_ok_mask, but that would affect the mean as well.
high_tolerance = 1e-6
testee_lkj = tfd.LKJ(dimension=dim,
concentration=concentration,
input_output_cholesky=input_output_cholesky,
validate_args=True)
x = testee_lkj.sample(num_samples, seed=seed)
importance_weights = (
tf.exp(-testee_lkj.log_prob(x)) *
_det_ok_mask(x, det_bounds, input_output_cholesky))
importance_maxima = (1. / det_bounds)**(concentration - 1) * tf.exp(
testee_lkj._log_normalization())
check1 = st.assert_true_mean_in_interval_by_dkwm(
samples=importance_weights,
low=0.,
high=importance_maxima + high_tolerance,
expected_low=lows,
expected_high=highs,
false_fail_rate=false_fail_rate)
check2 = assert_util.assert_less(
st.min_discrepancy_of_true_means_detectable_by_dkwm(
num_samples,
low=0.,
high=importance_maxima + high_tolerance,
false_fail_rate=false_fail_rate,
false_pass_rate=false_fail_rate), dtype(target_discrepancy))
self.evaluate([check1, check2])
def _parameter_control_dependencies(self, is_init):
assertions = []
axis = None
paddings = None
if is_init != tensor_util.is_ref(self.axis):
# First we check the shape of the axis argument.
msg = 'Argument `axis` must be scalar or vector.'
if tensorshape_util.rank(self.axis.shape) is not None:
if tensorshape_util.rank(self.axis.shape) > 1:
raise ValueError(msg)
elif self.validate_args:
if axis is None: axis = tf.convert_to_tensor(self.axis)
assertions.append(assert_util.assert_rank_at_most(
axis, 1, message=msg))
# Next we check the values of the axis argument.
axis_ = tf.get_static_value(self.axis)
msg = 'Argument `axis` must be negative.'
if axis_ is not None:
if np.any(axis_ > -1):
raise ValueError(msg)
elif self.validate_args:
if axis is None: axis = tf.convert_to_tensor(self.axis)
assertions.append(assert_util.assert_less(axis, 0, message=msg))
msg = 'Argument `axis` elements must be unique.'
if axis_ is not None:
if len(np.array(axis_).reshape(-1)) != len(np.unique(axis_)):
raise ValueError(msg)
elif self.validate_args:
if axis is None: axis = tf.convert_to_tensor(self.axis)
assertions.append(assert_util.assert_equal(
prefer_static.size0(axis),
prefer_static.size0(prefer_static.setdiff1d(axis)),
message=msg))
if is_init != tensor_util.is_ref(self.paddings):
# First we check the shape of the paddings argument.
msg = 'Argument `paddings` must be a vector of pairs.'
if tensorshape_util.is_fully_defined(self.paddings.shape):
shape = np.int32(self.paddings.shape)
if len(shape) != 2 or shape[0] < 1 or shape[1] != 2:
raise ValueError(msg)
elif self.validate_args:
if paddings is None: paddings = tf.convert_to_tensor(self.paddings)
with tf.control_dependencies([
assert_util.assert_equal(tf.rank(paddings), 2, message=msg)]):
shape = tf.shape(paddings)
assertions.extend([
assert_util.assert_greater(shape[0], 0, message=msg),
assert_util.assert_equal(shape[1], 2, message=msg),
])
# Next we check the values of the paddings argument.
paddings_ = tf.get_static_value(self.paddings)
msg = 'Argument `paddings` must be non-negative.'
if paddings_ is not None:
if np.any(paddings_ < 0):
raise ValueError(msg)
elif self.validate_args:
if paddings is None: paddings = tf.convert_to_tensor(self.paddings)
assertions.append(assert_util.assert_greater(
paddings, -1, message=msg))
if is_init != (tensor_util.is_ref(self.axis) and
tensor_util.is_ref(self.paddings)):
axis_ = tf.get_static_value(self.axis)
if axis_ is None and axis is None:
axis = tf.convert_to_tensor(self.axis)
len_axis = prefer_static.size0(prefer_static.reshape(
axis if axis_ is None else axis_, shape=-1))
paddings_ = tf.get_static_value(self.paddings)
if paddings_ is None and paddings is None:
paddings = tf.convert_to_tensor(self.paddings)
len_paddings = prefer_static.size0(
paddings if paddings_ is None else paddings_)
msg = ('Arguments `axis` and `paddings` must have the same number '
'of elements.')
if (prefer_static.is_numpy(len_axis) and
prefer_static.is_numpy(len_paddings)):
if len_axis != len_paddings:
raise ValueError(msg + ' Saw: {}, {}.'.format(
self.axis, self.paddings))
elif self.validate_args:
assertions.append(assert_util.assert_equal(
len_axis, len_paddings, message=msg))
return assertions