def testUniformUniformKLFinite(self):
batch_size = 6
a_low = -1.0 * np.arange(1, batch_size + 1)
a_high = np.array([1.0] * batch_size)
b_low = -2.0 * np.arange(1, batch_size + 1)
b_high = np.array([2.0] * batch_size)
a = uniform_lib.Uniform(low=a_low, high=a_high)
b = uniform_lib.Uniform(low=b_low, high=b_high)
true_kl = np.log(b_high - b_low) - np.log(a_high - a_low)
kl = tfd.kl_divergence(a, b)
# This is essentially an approximated integral from the direct definition
# of KL divergence.
x = a.sample(int(1e4), seed=0)
kl_sample = tf.reduce_mean(input_tensor=a.log_prob(x) - b.log_prob(x),
axis=0)
kl_, kl_sample_ = self.evaluate([kl, kl_sample])
self.assertAllEqual(true_kl, kl_)
self.assertAllClose(true_kl, kl_sample_, atol=0.0, rtol=1e-1)
zero_kl = tfd.kl_divergence(a, a)
true_zero_kl_, zero_kl_ = self.evaluate(
[tf.zeros_like(true_kl), zero_kl])
self.assertAllEqual(true_zero_kl_, zero_kl_)
def testFullyReparameterized(self):
a = tf.constant(0.1)
b = tf.constant(0.8)
_, [grad_a, grad_b] = tfp.math.value_and_gradient(
lambda a_, b_: uniform_lib.Uniform(a_, b_).sample(100), [a, b])
self.assertIsNotNone(grad_a)
self.assertIsNotNone(grad_b)
def _testUniformSampleMultiDimensional(self):
# DISABLED: Please enable this test once b/issues/30149644 is resolved.
batch_size = 2
a_v = [3.0, 22.0]
b_v = [13.0, 35.0]
a = tf.constant([a_v] * batch_size)
b = tf.constant([b_v] * batch_size)
uniform = uniform_lib.Uniform(low=a, high=b)
n_v = 100000
n = tf.constant(n_v)
samples = uniform.sample(n)
self.assertEqual(samples.shape, (n_v, batch_size, 2))
sample_values = self.evaluate(samples)
self.assertFalse(
np.any(sample_values[:, 0, 0] < a_v[0])
or np.any(sample_values[:, 0, 0] >= b_v[0]))
self.assertFalse(
np.any(sample_values[:, 0, 1] < a_v[1])
or np.any(sample_values[:, 0, 1] >= b_v[1]))
self.assertAllClose(sample_values[:, 0, 0].mean(),
(a_v[0] + b_v[0]) / 2,
atol=1e-2)
self.assertAllClose(sample_values[:, 0, 1].mean(),
(a_v[1] + b_v[1]) / 2,
atol=1e-2)
def __init__(self,
loc,
scale,
validate_args=False,
allow_nan_stats=True,
name="Gumbel"):
"""Construct Gumbel distributions with location and scale `loc` and `scale`.
The parameters `loc` and `scale` must be shaped in a way that supports
broadcasting (e.g. `loc + scale` is a valid operation).
Args:
loc: Floating point tensor, the means of the distribution(s).
scale: Floating point tensor, the scales of the distribution(s).
scale must contain only positive values.
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: `'Gumbel'`.
Raises:
TypeError: if loc and scale are different dtypes.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([loc, scale], dtype_hint=tf.float32)
loc = tf.convert_to_tensor(loc, name="loc", dtype=dtype)
scale = tf.convert_to_tensor(scale, name="scale", dtype=dtype)
with tf.control_dependencies(
[assert_util.assert_positive(scale)] if validate_args else []):
loc = tf.identity(loc, name="loc")
scale = tf.identity(scale, name="scale")
dtype_util.assert_same_float_dtype([loc, scale])
self._gumbel_bijector = gumbel_bijector.Gumbel(
loc=loc, scale=scale, validate_args=validate_args)
# Because the uniform sampler generates samples in `[0, 1)` this would
# cause samples to lie in `(inf, -inf]` instead of `(inf, -inf)`. To fix
# this, we use `np.finfo(dtype_util.as_numpy_dtype(self.dtype).tiny`
# because it is the smallest, positive, "normal" number.
super(Gumbel, self).__init__(
distribution=uniform.Uniform(
low=np.finfo(dtype_util.as_numpy_dtype(dtype)).tiny,
high=tf.ones([], dtype=loc.dtype),
allow_nan_stats=allow_nan_stats),
# The Gumbel bijector encodes the quantile
# function as the forward, and hence needs to
# be inverted.
bijector=invert_bijector.Invert(self._gumbel_bijector),
batch_shape=distribution_util.get_broadcast_shape(loc, scale),
parameters=parameters,
name=name)
def testUniformSamplePdf(self):
a = 10.0
b = [11.0, 100.0]
uniform = uniform_lib.Uniform(a, b)
self.assertTrue(
self.evaluate(
tf.reduce_all(uniform.prob(uniform.sample(10)) > 0)))
def testUniformRange(self):
a = 3.0
b = 10.0
uniform = uniform_lib.Uniform(low=a, high=b)
self.assertAllClose(a, self.evaluate(uniform.low))
self.assertAllClose(b, self.evaluate(uniform.high))
self.assertAllClose(b - a, self.evaluate(uniform.range()))
def testUniformEntropy(self):
a_v = np.array([1.0, 1.0, 1.0])
b_v = np.array([[1.5, 2.0, 3.0]])
uniform = uniform_lib.Uniform(low=a_v, high=b_v)
expected_entropy = np.log(b_v - a_v)
self.assertAllClose(expected_entropy, self.evaluate(uniform.entropy()))
def testUniformStd(self):
a = 10.0
b = 100.0
uniform = uniform_lib.Uniform(low=a, high=b)
if not stats:
return
s_uniform = stats.uniform(loc=a, scale=b - a)
self.assertAllClose(self.evaluate(uniform.stddev()), s_uniform.std())
def testUniformBroadcasting(self):
a = 10.0
b = [11.0, 20.0]
uniform = uniform_lib.Uniform(a, b)
pdf = uniform.prob([[10.5, 11.5], [9.0, 19.0], [10.5, 21.0]])
expected_pdf = np.array([[1.0, 0.1], [0.0, 0.1], [1.0, 0.0]])
self.assertAllClose(expected_pdf, self.evaluate(pdf))
def testUniformAssertMaxGtMin(self):
a_v = np.array([1.0, 1.0, 1.0], dtype=np.float32)
b_v = np.array([1.0, 2.0, 3.0], dtype=np.float32)
with self.assertRaisesWithPredicateMatch(errors.InvalidArgumentError,
"x < y"):
uniform = uniform_lib.Uniform(low=a_v, high=b_v, validate_args=True)
self.evaluate(uniform.low)
def __init__(self,
concentration,
scale,
validate_args=False,
allow_nan_stats=True,
name='Weibull'):
"""Construct Weibull distributions.
The parameters `concentration` and `scale` must be shaped in a way that
supports broadcasting (e.g. `concentration + scale` is a valid operation).
Args:
concentration: Positive Float-type `Tensor`, the concentration param of the
distribution. Must contain only positive values.
scale: Positive Float-type `Tensor`, the scale param of the distribution.
Must contain only positive values.
validate_args: Python `bool` indicating whether arguments should be checked
for correctness.
allow_nan_stats: Python `bool` indicating whether nan values should be
allowed.
name: Python `str` name given to ops managed by this class.
Default value: `'Weibull'`.
Raises:
TypeError: if concentration and scale are different dtypes.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([concentration, scale],
dtype_hint=tf.float32)
concentration = tensor_util.convert_nonref_to_tensor(
concentration, name='concentration', dtype=dtype)
scale = tensor_util.convert_nonref_to_tensor(scale,
name='scale',
dtype=dtype)
# Positive scale and concentration is asserted by the incorporated
# Weibull bijector.
self._weibull_bijector = weibull_cdf_bijector.WeibullCDF(
scale=scale,
concentration=concentration,
validate_args=validate_args)
batch_shape = distribution_util.get_broadcast_shape(
concentration, scale)
super(Weibull, self).__init__(
distribution=uniform.Uniform(
# TODO(b/137665504): Use batch-adding meta-distribution to set the
# batch shape instead of tf.ones.
low=tf.zeros(batch_shape, dtype=dtype),
high=tf.ones(batch_shape, dtype=dtype),
allow_nan_stats=allow_nan_stats),
# The Weibull bijector encodes the CDF function as the forward,
# and hence needs to be inverted.
bijector=invert_bijector.Invert(self._weibull_bijector,
validate_args=validate_args),
parameters=parameters,
name=name)
def testUniformShape(self):
a = tf.constant([-3.0] * 5)
b = tf.constant(11.0)
uniform = uniform_lib.Uniform(low=a, high=b)
self.assertEqual(self.evaluate(uniform.batch_shape_tensor()), (5, ))
self.assertEqual(uniform.batch_shape, tf.TensorShape([5]))
self.assertAllEqual(self.evaluate(uniform.event_shape_tensor()), [])
self.assertEqual(uniform.event_shape, tf.TensorShape([]))
def testUniformQuantile(self):
low = tf.reshape(tf.linspace(0., 1., 6), [2, 1, 3])
high = tf.reshape(tf.linspace(1.5, 2.5, 6), [1, 2, 3])
uniform = uniform_lib.Uniform(low=low, high=high)
expected_quantiles = tf.reshape(tf.linspace(1.01, 1.49, 24),
[2, 2, 2, 3])
cumulative_densities = uniform.cdf(expected_quantiles)
actual_quantiles = uniform.quantile(cumulative_densities)
self.assertAllClose(self.evaluate(expected_quantiles),
self.evaluate(actual_quantiles))
def testUniformPDFWithScalarEndpoint(self):
a = tf.constant([0.0, 5.0])
b = tf.constant(10.0)
uniform = uniform_lib.Uniform(low=a, high=b)
x = np.array([0.0, 8.0], dtype=np.float32)
expected_pdf = np.array([1.0 / (10.0 - 0.0), 1.0 / (10.0 - 5.0)])
pdf = uniform.prob(x)
self.assertAllClose(expected_pdf, self.evaluate(pdf))
def testFullyReparameterized(self):
a = tf.constant(0.1)
b = tf.constant(0.8)
with tf.GradientTape() as tape:
tape.watch(a)
tape.watch(b)
uniform = uniform_lib.Uniform(a, b)
samples = uniform.sample(100)
grad_a, grad_b = tape.gradient(samples, [a, b])
self.assertIsNotNone(grad_a)
self.assertIsNotNone(grad_b)
def testUniformUniformKLInfinite(self):
# This covers three cases:
# - a.low < b.low,
# - a.high > b.high, and
# - both.
a_low = np.array([-1.0, 0.0, -1.0])
a_high = np.array([1.0, 2.0, 2.0])
b_low = np.array([0.0] * 3)
b_high = np.array([1.0] * 3)
a = uniform_lib.Uniform(low=a_low, high=a_high)
b = uniform_lib.Uniform(low=b_low, high=b_high)
# Since 'a' can be sampled to give points outside the support of 'b',
# the KL Divergence is infinite.
true_kl = tf.convert_to_tensor(value=np.array([np.inf] * 3))
kl = tfd.kl_divergence(a, b)
true_kl_, kl_ = self.evaluate([true_kl, kl])
self.assertAllEqual(true_kl_, kl_)
def __init__(self,
concentration1=1.,
concentration0=1.,
validate_args=False,
allow_nan_stats=True,
name='Kumaraswamy'):
"""Initialize a batch of Kumaraswamy distributions.
Args:
concentration1: Positive floating-point `Tensor` indicating mean
number of successes; aka 'alpha'. Implies `self.dtype` and
`self.batch_shape`, i.e.,
`concentration1.shape = [N1, N2, ..., Nm] = self.batch_shape`.
concentration0: Positive floating-point `Tensor` indicating mean
number of failures; aka 'beta'. Otherwise has same semantics as
`concentration1`.
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.
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.
name: Python `str` name prefixed to Ops created by this class.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([concentration1, concentration0],
dtype_hint=tf.float32)
concentration1 = tensor_util.convert_nonref_to_tensor(
concentration1, name='concentration1', dtype=dtype)
concentration0 = tensor_util.convert_nonref_to_tensor(
concentration0, name='concentration0', dtype=dtype)
self._kumaraswamy_cdf = kumaraswamy_cdf.KumaraswamyCDF(
concentration1=concentration1,
concentration0=concentration0,
validate_args=validate_args)
batch_shape = distribution_util.get_broadcast_shape(
concentration1, concentration0)
super(Kumaraswamy, self).__init__(
# TODO(b/137665504): Use batch-adding meta-distribution to set the
# batch shape instead of tf.zeros.
distribution=uniform.Uniform(
low=tf.zeros(batch_shape, dtype=dtype),
high=tf.ones([], dtype=dtype),
allow_nan_stats=allow_nan_stats),
bijector=invert.Invert(
self._kumaraswamy_cdf, validate_args=validate_args),
parameters=parameters,
name=name)
def __init__(self,
concentration1=None,
concentration0=None,
validate_args=False,
allow_nan_stats=True,
name="Kumaraswamy"):
"""Initialize a batch of Kumaraswamy distributions.
Args:
concentration1: Positive floating-point `Tensor` indicating mean
number of successes; aka "alpha". Implies `self.dtype` and
`self.batch_shape`, i.e.,
`concentration1.shape = [N1, N2, ..., Nm] = self.batch_shape`.
concentration0: Positive floating-point `Tensor` indicating mean
number of failures; aka "beta". Otherwise has same semantics as
`concentration1`.
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.
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.
name: Python `str` name prefixed to Ops created by this class.
"""
parameters = dict(locals())
with tf.compat.v1.name_scope(name,
values=[concentration1,
concentration0]) as name:
dtype = dtype_util.common_dtype([concentration1, concentration0],
tf.float32)
concentration1 = tf.convert_to_tensor(value=concentration1,
name="concentration1",
dtype=dtype)
concentration0 = tf.convert_to_tensor(value=concentration0,
name="concentration0",
dtype=dtype)
super(Kumaraswamy,
self).__init__(distribution=uniform.Uniform(
low=tf.zeros([], dtype=concentration1.dtype),
high=tf.ones([], dtype=concentration1.dtype),
allow_nan_stats=allow_nan_stats),
bijector=kumaraswamy_bijector.Kumaraswamy(
concentration1=concentration1,
concentration0=concentration0,
validate_args=validate_args),
batch_shape=distribution_util.get_broadcast_shape(
concentration1, concentration0),
parameters=parameters,
name=name)
def testUniformNans(self):
a = 10.0
b = [11.0, 100.0]
uniform = uniform_lib.Uniform(low=a, high=b)
no_nans = tf.constant(1.0)
nans = tf.constant(0.0) / tf.constant(0.0)
self.assertTrue(self.evaluate(tf.is_nan(nans)))
with_nans = tf.stack([no_nans, nans])
pdf = uniform.prob(with_nans)
is_nan = self.evaluate(tf.is_nan(pdf))
self.assertFalse(is_nan[0])
self.assertTrue(is_nan[1])
def testUniformFloat64(self):
uniform = uniform_lib.Uniform(low=np.float64(0.), high=np.float64(1.))
self.assertAllClose(
[1., 1.],
self.evaluate(uniform.prob(np.array([0.5, 0.6],
dtype=np.float64))))
self.assertAllClose(
[0.5, 0.6],
self.evaluate(uniform.cdf(np.array([0.5, 0.6], dtype=np.float64))))
self.assertAllClose(0.5, self.evaluate(uniform.mean()))
self.assertAllClose(1 / 12., self.evaluate(uniform.variance()))
self.assertAllClose(0., self.evaluate(uniform.entropy()))
def testUniformSampleWithShape(self):
a = 10.0
b = [11.0, 20.0]
uniform = uniform_lib.Uniform(a, b)
pdf = uniform.prob(uniform.sample((2, 3)))
# pylint: disable=bad-continuation
expected_pdf = [
[[1.0, 0.1], [1.0, 0.1], [1.0, 0.1]],
[[1.0, 0.1], [1.0, 0.1], [1.0, 0.1]],
]
# pylint: enable=bad-continuation
self.assertAllClose(expected_pdf, self.evaluate(pdf))
pdf = uniform.prob(uniform.sample())
expected_pdf = [1.0, 0.1]
self.assertAllClose(expected_pdf, self.evaluate(pdf))
def testUniformSample(self):
a = tf.constant([3.0, 4.0])
b = tf.constant(13.0)
a1_v = 3.0
a2_v = 4.0
b_v = 13.0
n = tf.constant(100000)
uniform = uniform_lib.Uniform(low=a, high=b)
samples = uniform.sample(n, seed=137)
sample_values = self.evaluate(samples)
self.assertEqual(sample_values.shape, (100000, 2))
self.assertAllClose(
sample_values[::, 0].mean(), (b_v + a1_v) / 2, atol=1e-1, rtol=0.)
self.assertAllClose(
sample_values[::, 1].mean(), (b_v + a2_v) / 2, atol=1e-1, rtol=0.)
self.assertFalse(
np.any(sample_values[::, 0] < a1_v) or np.any(sample_values >= b_v))
self.assertFalse(
np.any(sample_values[::, 1] < a2_v) or np.any(sample_values >= b_v))
def testUniformCDF(self):
batch_size = 6
a = tf.constant([1.0] * batch_size)
b = tf.constant([11.0] * batch_size)
a_v = 1.0
b_v = 11.0
x = np.array([-2.5, 2.5, 4.0, 0.0, 10.99, 12.0], dtype=np.float32)
uniform = uniform_lib.Uniform(low=a, high=b)
def _expected_cdf():
cdf = (x - a_v) / (b_v - a_v)
cdf[x >= b_v] = 1
cdf[x < a_v] = 0
return cdf
cdf = uniform.cdf(x)
self.assertAllClose(_expected_cdf(), self.evaluate(cdf))
log_cdf = uniform.log_cdf(x)
self.assertAllClose(np.log(_expected_cdf()), self.evaluate(log_cdf))
def _uniform_correlation_like_matrix(num_rows, batch_shape, dtype, seed):
"""Returns a uniformly random `Tensor` of "correlation-like" matrices.
A "correlation-like" matrix is a symmetric square matrix with all entries
between -1 and 1 (inclusive) and 1s on the main diagonal. Of these,
the ones that are positive semi-definite are exactly the correlation
matrices.
Args:
num_rows: Python `int` dimension of the correlation-like matrices.
batch_shape: `Tensor` or Python `tuple` of `int` shape of the
batch to return.
dtype: `dtype` of the `Tensor` to return.
seed: Random seed.
Returns:
matrices: A `Tensor` of shape `batch_shape + [num_rows, num_rows]`
and dtype `dtype`. Each entry is in [-1, 1], and each matrix
along the bottom two dimensions is symmetric and has 1s on the
main diagonal.
"""
num_entries = num_rows * (num_rows + 1) / 2
ones = tf.ones(shape=[num_entries], dtype=dtype)
# It seems wasteful to generate random values for the diagonal since
# I am going to throw them away, but `fill_triangular` fills the
# diagonal, so I probably need them.
# It's not impossible that it would be more efficient to just fill
# the whole matrix with random values instead of messing with
# `fill_triangular`. Then would need to filter almost half out with
# `matrix_band_part`.
unifs = uniform.Uniform(-ones, ones).sample(batch_shape, seed=seed)
tril = util.fill_triangular(unifs)
symmetric = tril + tf.linalg.transpose(tril)
diagonal_ones = tf.ones(shape=util.pad(batch_shape,
axis=0,
back=True,
value=num_rows),
dtype=dtype)
return tf.linalg.set_diag(symmetric, diagonal_ones)
def testUniformPDF(self):
a = tf.constant([-3.0] * 5 + [15.0])
b = tf.constant([11.0] * 5 + [20.0])
uniform = uniform_lib.Uniform(low=a, high=b)
a_v = -3.0
b_v = 11.0
x = np.array([-10.5, 4.0, 0.0, 10.99, 11.3, 17.0], dtype=np.float32)
def _expected_pdf():
pdf = np.zeros_like(x) + 1.0 / (b_v - a_v)
pdf[x > b_v] = 0.0
pdf[x < a_v] = 0.0
pdf[5] = 1.0 / (20.0 - 15.0)
return pdf
expected_pdf = _expected_pdf()
pdf = uniform.prob(x)
self.assertAllClose(expected_pdf, self.evaluate(pdf))
log_pdf = uniform.log_prob(x)
self.assertAllClose(np.log(expected_pdf), self.evaluate(log_pdf))
def resample_minimum_variance(log_probs,
event_size,
sample_shape,
seed=None,
name=None):
"""Minimum variance resampler for sequential Monte Carlo.
This function is based on Algorithm #2 in [Maskell et al. (2006)][1].
Args:
log_probs: A tensor-valued batch of discrete log probability distributions.
event_size: the dimension of the vector considered a single draw.
sample_shape: the `sample_shape` determining the number of draws.
seed: Python '`int` used to seed calls to `tf.random.*`.
Default value: None (i.e. no seed).
name: Python `str` name for ops created by this method.
Default value: `None` (i.e., `'resample_minimum_variance'`).
Returns:
resampled_indices: The result is similar to sampling with
```python
expanded_sample_shape = tf.concat([[event_size], sample_shape]), axis=-1)
tfd.Categorical(logits=log_probs).sample(expanded_sample_shape)`
```
but with values sorted along the first axis. It can be considered to be
sampling events made up of a length-`event_size` vector of draws from
the `Categorical` distribution. However, although the elements of
this event have the appropriate marginal distribution, they are not
independent of each other. Instead they have been chosen so as to form
a good representative sample, suitable for use with Sequential Monte
Carlo algorithms.
The sortedness is an unintended side effect of the algorithm that is
harmless in the context of simple SMC algorithms.
#### References
[1]: S. Maskell, B. Alun-Jones and M. Macleod. A Single Instruction Multiple
Data Particle Filter.
In 2006 IEEE Nonlinear Statistical Signal Processing Workshop.
http://people.ds.cam.ac.uk/fanf2/hermes/doc/antiforgery/stats.pdf
"""
with tf.name_scope(name or 'resample_minimum_variance') as name:
log_probs = tf.convert_to_tensor(log_probs, dtype_hint=tf.float32)
log_probs = dist_util.move_dimension(log_probs,
source_idx=0,
dest_idx=-1)
batch_shape = prefer_static.shape(log_probs)[:-1]
working_shape = prefer_static.concat([sample_shape, batch_shape],
axis=-1)
log_cdf = tf.math.cumulative_logsumexp(log_probs[..., :-1], axis=-1)
# Each resampling requires a single uniform random variable
offset = uniform.Uniform(low=tf.constant(0., log_cdf.dtype),
high=tf.constant(1., log_cdf.dtype)).sample(
working_shape, seed=seed)[..., tf.newaxis]
# It is possible for numerical error to result in a cumulative
# sum that exceeds 1 so we need to clip.
markers = prefer_static.cast(
tf.floor(event_size * tf.math.exp(log_cdf) + offset), tf.int32)
indices = markers[..., tf.newaxis]
updates = tf.ones(prefer_static.shape(indices)[:-1], dtype=tf.int32)
scatter_shape = prefer_static.concat([working_shape, [event_size + 1]],
axis=-1)
batch_dims = (prefer_static.rank_from_shape(sample_shape) +
prefer_static.rank_from_shape(batch_shape))
x = _scatter_nd_batch(indices,
updates,
scatter_shape,
batch_dims=batch_dims)
resampled = tf.cumsum(x, axis=-1)[..., :-1]
resampled = dist_util.move_dimension(resampled,
source_idx=-1,
dest_idx=0)
return resampled
def __init__(self,
loc,
scale,
validate_args=False,
allow_nan_stats=True,
name='Moyal'):
"""Construct Moyal distributions with location and scale `loc` and `scale`.
The parameters `loc` and `scale` must be shaped in a way that supports
broadcasting (e.g. `loc + scale` is a valid operation).
Args:
loc: Floating point tensor, the means of the distribution(s).
scale: Floating point tensor, the scales of the distribution(s).
scale must contain only positive values.
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: `'Moyal'`.
Raises:
TypeError: if loc and scale are different dtypes.
#### References
[1] J.E. Moyal, "XXX. Theory of ionization fluctuations",
The London, Edinburgh, and Dublin Philosophical Magazine
and Journal of Science.
https://www.tandfonline.com/doi/abs/10.1080/14786440308521076
[2] G. Cordeiro, J. Nobre, R. Pescim, E. Ortega,
"The beta Moyal: a useful skew distribution",
https://www.arpapress.com/Volumes/Vol10Issue2/IJRRAS_10_2_02.pdf
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([loc, scale],
dtype_hint=tf.float32)
loc = tensor_util.convert_nonref_to_tensor(loc,
name='loc',
dtype=dtype)
scale = tensor_util.convert_nonref_to_tensor(scale,
name='scale',
dtype=dtype)
dtype_util.assert_same_float_dtype([loc, scale])
# Positive scale is asserted by the incorporated Moyal bijector.
self._moyal_bijector = moyal_cdf_bijector.MoyalCDF(
loc=loc, scale=scale, validate_args=validate_args)
# Because the uniform sampler generates samples in `[0, 1)` this would
# cause samples to lie in `(inf, -inf]` instead of `(inf, -inf)`. To fix
# this, we use `np.finfo(dtype_util.as_numpy_dtype(self.dtype).tiny`
# because it is the smallest, positive, 'normal' number.
batch_shape = distribution_util.get_broadcast_shape(loc, scale)
super(Moyal, self).__init__(
# TODO(b/137665504): Use batch-adding meta-distribution to set the
# batch shape instead of tf.ones.
distribution=uniform.Uniform(low=np.finfo(
dtype_util.as_numpy_dtype(dtype)).tiny,
high=tf.ones(batch_shape,
dtype=dtype),
allow_nan_stats=allow_nan_stats),
# The Moyal bijector encodes the CDF function as the forward,
# and hence needs to be inverted.
bijector=invert_bijector.Invert(self._moyal_bijector,
validate_args=validate_args),
parameters=parameters,
name=name)
def assert_scalar_congruency(bijector,
lower_x,
upper_x,
eval_func,
n=int(10e3),
rtol=0.01):
"""Assert `bijector`'s forward/inverse/inverse_log_det_jacobian are congruent.
We draw samples `X ~ U(lower_x, upper_x)`, then feed these through the
`bijector` in order to check that:
1. the forward is strictly monotonic.
2. the forward/inverse methods are inverses of each other.
3. the jacobian is the correct change of measure.
This can only be used for a Bijector mapping open subsets of the real line
to themselves. This is due to the fact that this test compares the `prob`
before/after transformation with the Lebesgue measure on the line.
Args:
bijector: Instance of Bijector
lower_x: Python scalar.
upper_x: Python scalar. Must have `lower_x < upper_x`, and both must be in
the domain of the `bijector`. The `bijector` should probably not produce
huge variation in values in the interval `(lower_x, upper_x)`, or else the
variance based check of the Jacobian will require small `rtol` or huge
`n`.
eval_func: Function to evaluate any intermediate results.
n: Number of samples to draw for the checks.
rtol: Positive number. Used for the Jacobian check.
Raises:
AssertionError: If tests fail.
"""
# Should be monotonic over this interval
ten_x_pts = np.linspace(lower_x, upper_x, num=10).astype(np.float32)
if bijector.dtype is not None:
ten_x_pts = ten_x_pts.astype(dtype_util.as_numpy_dtype(bijector.dtype))
forward_on_10_pts = bijector.forward(ten_x_pts)
# Set the lower/upper limits in the range of the bijector.
lower_y, upper_y = eval_func(
[bijector.forward(lower_x),
bijector.forward(upper_x)])
if upper_y < lower_y: # If bijector.forward is a decreasing function.
lower_y, upper_y = upper_y, lower_y
# Uniform samples from the domain, range.
uniform_x_samps = uniform_distribution.Uniform(low=lower_x,
high=upper_x).sample(n,
seed=0)
uniform_y_samps = uniform_distribution.Uniform(low=lower_y,
high=upper_y).sample(n,
seed=1)
# These compositions should be the identity.
inverse_forward_x = bijector.inverse(bijector.forward(uniform_x_samps))
forward_inverse_y = bijector.forward(bijector.inverse(uniform_y_samps))
# For a < b, and transformation y = y(x),
# (b - a) = \int_a^b dx = \int_{y(a)}^{y(b)} |dx/dy| dy
# "change_measure_dy_dx" below is a Monte Carlo approximation to the right
# hand side, which should then be close to the left, which is (b - a).
# We assume event_ndims=0 because we assume scalar -> scalar. The log_det
# methods will handle whether they expect event_ndims > 0.
dy_dx = tf.exp(
bijector.inverse_log_det_jacobian(uniform_y_samps, event_ndims=0))
# E[|dx/dy|] under Uniform[lower_y, upper_y]
# = \int_{y(a)}^{y(b)} |dx/dy| dP(u), where dP(u) is the uniform measure
expectation_of_dy_dx_under_uniform = tf.reduce_mean(dy_dx)
# dy = dP(u) * (upper_y - lower_y)
change_measure_dy_dx = ((upper_y - lower_y) *
expectation_of_dy_dx_under_uniform)
# We'll also check that dy_dx = 1 / dx_dy.
dx_dy = tf.exp(
bijector.forward_log_det_jacobian(bijector.inverse(uniform_y_samps),
event_ndims=0))
[
forward_on_10_pts_v,
dy_dx_v,
dx_dy_v,
change_measure_dy_dx_v,
uniform_x_samps_v,
uniform_y_samps_v,
inverse_forward_x_v,
forward_inverse_y_v,
] = eval_func([
forward_on_10_pts,
dy_dx,
dx_dy,
change_measure_dy_dx,
uniform_x_samps,
uniform_y_samps,
inverse_forward_x,
forward_inverse_y,
])
assert_strictly_monotonic(forward_on_10_pts_v)
# Composition of forward/inverse should be the identity.
np.testing.assert_allclose(inverse_forward_x_v,
uniform_x_samps_v,
atol=1e-5,
rtol=1e-3)
np.testing.assert_allclose(forward_inverse_y_v,
uniform_y_samps_v,
atol=1e-5,
rtol=1e-3)
# Change of measure should be correct.
np.testing.assert_allclose(upper_x - lower_x,
change_measure_dy_dx_v,
atol=0,
rtol=rtol)
# Inverse Jacobian should be equivalent to the reciprocal of the forward
# Jacobian.
np.testing.assert_allclose(dy_dx_v,
np.divide(1., dx_dy_v),
atol=1e-5,
rtol=1e-3)
def resample_stratified(log_probs,
event_size,
sample_shape,
seed=None,
name=None):
"""Stratified resampler for sequential Monte Carlo.
The value returned from this algorithm is similar to sampling with
```python
expanded_sample_shape = tf.concat([[event_size], sample_shape]), axis=-1)
tfd.Categorical(logits=log_probs).sample(expanded_sample_shape)`
```
but with values sorted along the first axis. It can be considered to be
sampling events made up of a length-`event_size` vector of draws from
the `Categorical` distribution. However, although the elements of
this event have the appropriate marginal distribution, they are not
independent of each other. Instead they are drawn using a low variance
stratified sampling method suitable for use with Sequential Monte
Carlo algorithms.
The sortedness is an unintended side effect of the algorithm that is
harmless in the context of simple SMC algorithms.
This function is based on Algorithm #1 in the paper
[Maskell et al. (2006)][1].
Args:
log_probs: A tensor-valued batch of discrete log probability distributions.
event_size: the dimension of the vector considered a single draw.
sample_shape: the `sample_shape` determining the number of draws.
seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
Default value: None (i.e. no seed).
name: Python `str` name for ops created by this method.
Default value: `None` (i.e., `'resample_independent'`).
Returns:
resampled_indices: a tensor of samples.
#### References
[1]: S. Maskell, B. Alun-Jones and M. Macleod. A Single Instruction Multiple
Data Particle Filter.
In 2006 IEEE Nonlinear Statistical Signal Processing Workshop.
http://people.ds.cam.ac.uk/fanf2/hermes/doc/antiforgery/stats.pdf
"""
with tf.name_scope(name or 'resample_stratified') as name:
log_probs = tf.convert_to_tensor(log_probs, dtype_hint=tf.float32)
log_probs = dist_util.move_dimension(log_probs,
source_idx=0,
dest_idx=-1)
points_shape = ps.concat(
[sample_shape,
ps.shape(log_probs)[:-1], [event_size]], axis=0)
# Draw an offset for every element of an event.
interval_width = ps.cast(1. / event_size, dtype=log_probs.dtype)
offsets = uniform.Uniform(low=ps.cast(0., dtype=log_probs.dtype),
high=interval_width).sample(points_shape,
seed=seed)
# The unit interval is divided into equal partitions and each point
# is a random offset into a partition.
even_spacing = tf.linspace(start=ps.cast(0., dtype=log_probs.dtype),
stop=1 - interval_width,
num=event_size) + offsets
log_points = tf.math.log(even_spacing)
resampled = _resample_using_log_points(log_probs, sample_shape,
log_points)
return dist_util.move_dimension(resampled, source_idx=-1, dest_idx=0)
def resample_systematic(log_probs,
event_size,
sample_shape,
seed=None,
name=None):
"""A systematic resampler for sequential Monte Carlo.
The value returned from this function is similar to sampling with
```python
expanded_sample_shape = tf.concat([[event_size], sample_shape]), axis=-1)
tfd.Categorical(logits=log_probs).sample(expanded_sample_shape)`
```
but with values sorted along the first axis. It can be considered to be
sampling events made up of a length-`event_size` vector of draws from
the `Categorical` distribution. However, although the elements of
this event have the appropriate marginal distribution, they are not
independent of each other. Instead they are drawn using a stratified
sampling method that in some sense reduces variance and is suitable for
use with Sequential Monte Carlo algorithms as described in
[Doucet et al. (2011)][2].
The sortedness is an unintended side effect of the algorithm that is
harmless in the context of simple SMC algorithms.
This implementation is based on the algorithms in [Maskell et al. (2006)][1]
where it is called minimum variance resampling.
Args:
log_probs: A tensor-valued batch of discrete log probability distributions.
event_size: the dimension of the vector considered a single draw.
sample_shape: the `sample_shape` determining the number of draws.
seed: PRNG seed; see `tfp.random.sanitize_seed` for details.
Default value: None (i.e. no seed).
name: Python `str` name for ops created by this method.
Default value: `None` (i.e., `'resample_systematic'`).
Returns:
resampled_indices: a tensor of samples.
#### References
[1]: S. Maskell, B. Alun-Jones and M. Macleod. A Single Instruction Multiple
Data Particle Filter.
In 2006 IEEE Nonlinear Statistical Signal Processing Workshop.
http://people.ds.cam.ac.uk/fanf2/hermes/doc/antiforgery/stats.pdf
[2]: A. Doucet & A. M. Johansen. Tutorial on Particle Filtering and
Smoothing: Fifteen Years Later
In 2011 The Oxford Handbook of Nonlinear Filtering
https://www.stats.ox.ac.uk/~doucet/doucet_johansen_tutorialPF2011.pdf
"""
with tf.name_scope(name or 'resample_systematic') as name:
log_probs = tf.convert_to_tensor(log_probs, dtype_hint=tf.float32)
log_probs = dist_util.move_dimension(log_probs,
source_idx=0,
dest_idx=-1)
working_shape = ps.concat(
[sample_shape, ps.shape(log_probs)[:-1]], axis=0)
points_shape = ps.concat([working_shape, [event_size]], axis=0)
# Draw a single offset for each event.
interval_width = ps.cast(1. / event_size, dtype=log_probs.dtype)
offsets = uniform.Uniform(low=ps.cast(0., dtype=log_probs.dtype),
high=interval_width).sample(
working_shape, seed=seed)[...,
tf.newaxis]
even_spacing = ps.linspace(start=ps.cast(0., dtype=log_probs.dtype),
stop=1 - interval_width,
num=event_size) + offsets
log_points = tf.broadcast_to(tf.math.log(even_spacing), points_shape)
resampled = _resample_using_log_points(log_probs, sample_shape,
log_points)
return dist_util.move_dimension(resampled, source_idx=-1, dest_idx=0)
