def _sample_n(self, n, seed):
batch_shape = self.batch_shape_tensor()
event_shape = self.event_shape_tensor()
batch_ndims = tf.shape(batch_shape)[0]
ndims = batch_ndims + 3 # sample_ndims=1, event_ndims=2
shape = tf.concat([[n], batch_shape, event_shape], 0)
stream = seed_stream.SeedStream(seed, salt="Wishart")
# Complexity: O(nbk**2)
x = tf.random_normal(shape=shape,
mean=0.,
stddev=1.,
dtype=self.dtype,
seed=stream())
# Complexity: O(nbk)
# This parametrization is equivalent to Chi2, i.e.,
# ChiSquared(k) == Gamma(alpha=k/2, beta=1/2)
expanded_df = self.df * tf.ones(
self.scale_operator.batch_shape_tensor(),
dtype=self.df.dtype.base_dtype)
g = tf.random_gamma(shape=[n],
alpha=self._multi_gamma_sequence(
0.5 * expanded_df, self.dimension),
beta=0.5,
dtype=self.dtype,
seed=stream())
# Complexity: O(nbk**2)
x = tf.matrix_band_part(x, -1, 0) # Tri-lower.
# Complexity: O(nbk)
x = tf.matrix_set_diag(x, tf.sqrt(g))
# Make batch-op ready.
# Complexity: O(nbk**2)
perm = tf.concat([tf.range(1, ndims), [0]], 0)
x = tf.transpose(x, perm)
shape = tf.concat([batch_shape, [event_shape[0]], [-1]], 0)
x = tf.reshape(x, shape)
# Complexity: O(nbM) where M is the complexity of the operator solving a
# vector system. For LinearOperatorLowerTriangular, each matmul is O(k^3) so
# this step has complexity O(nbk^3).
x = self.scale_operator.matmul(x)
# Undo make batch-op ready.
# Complexity: O(nbk**2)
shape = tf.concat([batch_shape, event_shape, [n]], 0)
x = tf.reshape(x, shape)
perm = tf.concat([[ndims - 1], tf.range(0, ndims - 1)], 0)
x = tf.transpose(x, perm)
if not self.input_output_cholesky:
# Complexity: O(nbk**3)
x = tf.matmul(x, x, adjoint_b=True)
return x
def _sample_n(self, n, seed=None):
low = tf.convert_to_tensor(self.low)
high = tf.convert_to_tensor(self.high)
peak = tf.convert_to_tensor(self.peak)
stream = seed_stream.SeedStream(seed, salt='triangular')
shape = tf.concat(
[[n], self._batch_shape_tensor(low=low, high=high, peak=peak)],
axis=0)
samples = tf.random.uniform(shape=shape,
dtype=self.dtype,
seed=stream())
# We use Inverse CDF sampling here. Because the CDF is a quadratic function,
# we must use sqrts here.
interval_length = high - low
return tf.where(
# Note the CDF on the left side of the peak is
# (x - low) ** 2 / ((high - low) * (peak - low)).
# If we plug in peak for x, we get that the CDF at the peak
# is (peak - low) / (high - low). Because of this we decide
# which part of the piecewise CDF we should use based on the cdf samples
# we drew.
samples < (peak - low) / interval_length,
# Inverse of (x - low) ** 2 / ((high - low) * (peak - low)).
low + tf.sqrt(samples * interval_length * (peak - low)),
# Inverse of 1 - (high - x) ** 2 / ((high - low) * (high - peak))
high - tf.sqrt((1. - samples) * interval_length * (high - peak)))
def _sample_3d(self, n, seed=None):
"""Specialized inversion sampler for 3D."""
seed = seed_stream.SeedStream(seed, salt='von_mises_fisher_3d')
u_shape = tf.concat([[n], self._batch_shape_tensor()], axis=0)
z = tf.random.uniform(u_shape, seed=seed(), dtype=self.dtype)
# TODO(bjp): Higher-order odd dim analytic CDFs are available in [1], could
# be bisected for bounded sampling runtime (i.e. not rejection sampling).
# [1]: Inversion sampler via: https://ieeexplore.ieee.org/document/7347705/
# The inversion is: u = 1 + log(z + (1-z)*exp(-2*kappa)) / kappa
# We must protect against both kappa and z being zero.
safe_conc = tf.where(self.concentration > 0, self.concentration,
tf.ones_like(self.concentration))
safe_z = tf.where(z > 0, z, tf.ones_like(z))
safe_u = 1 + tf.reduce_logsumexp(
[tf.math.log(safe_z),
tf.math.log1p(-safe_z) - 2 * safe_conc],
axis=0) / safe_conc
# Limit of the above expression as kappa->0 is 2*z-1
u = tf.where(self.concentration > tf.zeros_like(safe_u), safe_u,
2 * z - 1)
# Limit of the expression as z->0 is -1.
u = tf.where(tf.equal(z, 0), -tf.ones_like(u), u)
if not self._allow_nan_stats:
u = tf.debugging.check_numerics(u, 'u in _sample_3d')
return u[..., tf.newaxis]
def _sample_n(self, n, seed=None):
# See https://en.wikipedia.org/wiki/Inverse_Gaussian_distribution or
# https://www.jstor.org/stable/2683801
concentration = tf.convert_to_tensor(self.concentration)
loc = tf.convert_to_tensor(self.loc)
seed = seed_stream.SeedStream(seed, 'inverse_gaussian')
shape = tf.concat(
[[n],
self._batch_shape_tensor(loc=loc, concentration=concentration)],
axis=0)
sampled_chi2 = (tf.random.normal(shape,
mean=0.,
stddev=1.,
seed=seed(),
dtype=self.dtype))**2.
sampled_uniform = tf.random.uniform(shape,
minval=0.,
maxval=1.,
seed=seed(),
dtype=self.dtype)
sampled = (loc + loc**2. * sampled_chi2 / (2. * concentration) - loc /
(2. * concentration) *
(4. * loc * concentration * sampled_chi2 +
(loc * sampled_chi2)**2)**0.5)
return tf.where(sampled_uniform <= loc / (loc + sampled), sampled,
loc**2 / sampled)
def _sample_n(self, n, seed=None):
seed = seed_stream.SeedStream(seed, salt="Autoregressive")()
samples = self.distribution0.sample(n, seed=seed)
for _ in range(self._num_steps):
# pylint: disable=not-callable
samples = self.distribution_fn(samples).sample(seed=seed)
return samples
def _flat_sample_distributions(self, sample_shape=(), seed=None, value=None):
"""Executes `model`, creating both samples and distributions."""
ds = []
values_out = []
seed = seed_stream.SeedStream('JointDistributionCoroutine', seed)
gen = self._model()
index = 0
d = next(gen)
try:
while True:
actual_distribution = d.distribution if isinstance(d, self.Root) else d
ds.append(actual_distribution)
if (value is not None and len(value) > index and
value[index] is not None):
seed()
next_value = value[index]
else:
next_value = actual_distribution.sample(
sample_shape=sample_shape if isinstance(d, self.Root) else (),
seed=seed())
values_out.append(next_value)
index += 1
d = gen.send(next_value)
except StopIteration:
pass
return ds, values_out
def _sample_n(self, n, seed=None):
# This implementation is equivalent to the one in tf.contrib.distributions.Wishart
batch_shape = self.batch_shape_tensor()
event_shape = self.event_shape_tensor()
stream = seed_stream.SeedStream(seed=seed, salt="Wishart")
shape = tf.concat([[n], batch_shape, event_shape], 0)
# Sample a normal full matrix
x = tf.random_normal(shape=shape, dtype=self.dtype, seed=stream())
# Sample the diagonal
g = tf.random_gamma(shape=[n],
alpha=self._multi_gamma_sequence(
0.5 * self.df, self.p),
beta=0.5,
dtype=self.dtype,
seed=stream())
# Discard the upper triangular part
x = tf.matrix_band_part(x, -1, 0)
# Set the diagonal
x = tf.matrix_set_diag(x, tf.sqrt(g))
# Scale with the Scale matrix, equivalent to matmul(sqrt(diag_scale), x)
x *= tf.sqrt(tf.exp(self.log_diag_scale[tf.newaxis, :, :, tf.newaxis]))
return x
def _flat_sample_distributions(self,
sample_shape=(),
seed=None,
value=None):
# This function additionally depends on:
# self._dist_fn_wrapped
# self._dist_fn_args
# self._always_use_specified_sample_shape
seed = seed_stream.SeedStream('JointDistributionSequential', seed)
ds = []
xs = [None] * len(self._dist_fn_wrapped) if value is None else list(
value)
if len(xs) != len(self._dist_fn_wrapped):
raise ValueError('Number of `xs`s must match number of '
'distributions.')
for i, (dist_fn, args) in enumerate(
zip(self._dist_fn_wrapped, self._dist_fn_args)):
ds.append(dist_fn(*xs[:i])) # Chain rule of probability.
if xs[i] is None:
# TODO(b/129364796): We should ignore args prefixed with `_`; this
# would mean we more often identify when to use `sample_shape=()`
# rather than `sample_shape=sample_shape`.
xs[i] = ds[-1].sample(
() if args and not self._always_use_specified_sample_shape
else sample_shape,
seed=seed())
else:
xs[i] = tf.convert_to_tensor(value=xs[i],
dtype_hint=ds[-1].dtype)
seed(
) # Ensure reproducibility even when xs are (partially) set.
# Note: we could also resolve distributions up to the first non-`None` in
# `self._flatten(value)`, however we omit this feature for simplicity,
# speed, and because it has not yet been requested.
return ds, xs
def _sample_distributions(self, sample_shape=(), seed=None, value=None):
seed = seed_stream.SeedStream('JointDistributionSequential', seed)
ds = []
if value is None:
xs = [None] * len(self._wrapped_dist_fn)
else:
xs = list(value)
xs.extend([None] * (len(self._wrapped_dist_fn) - len(xs)))
if len(xs) != len(self.distribution_fn):
raise ValueError('Number of `xs`s must match number of '
'distributions.')
for i, (dist_fn, args) in enumerate(
zip(self._wrapped_dist_fn, self._dist_fn_args)):
ds.append(dist_fn(*xs[:i])) # Chain rule of probability.
if xs[i] is None:
# TODO(b/129364796): We should ignore args prefixed with `_`; this
# would mean we more often identify when to use `sample_shape=()`
# rather than `sample_shape=sample_shape`.
xs[i] = ds[-1].sample(
() if args and not self._always_use_specified_sample_shape
else sample_shape,
seed=seed())
else:
xs[i] = tf.convert_to_tensor(value=xs[i],
dtype_hint=ds[-1].dtype)
seed(
) # Ensure reproducibility even when xs are (partially) set.
self._most_recently_built_distributions = ds
return tuple(ds), tuple(xs)
def _sample_n(self, n, seed=None):
with tf.control_dependencies(self._assertions):
seed = seed_stream.SeedStream(seed=seed,
salt='BlockwiseDistribution')
xs = [
tf.cast(d.sample(n, seed=seed()), self.dtype)
for d in self.distributions
]
return tf.concat(xs, axis=-1)
def _sample_n(self, n, seed):
with tf.compat.v1.control_dependencies(self._runtime_assertions):
seed = seed_stream.SeedStream(seed, salt="ZeroInflated")
mask = self.inflated_distribution.sample(n, seed())
samples = self.count_distribution.sample(n, seed())
mask, samples = _broadcast_rate(mask, samples)
# mask = 1 => new_sample = 0
# mask = 0 => new_sample = sample
return samples * tf.cast(1 - mask, samples.dtype)
def _sample_n_sparse(self, n, kw, seed=None):
assert n == 1
batch_shape = self.batch_shape_tensor()
event_shape = self.event_shape
iw = int(np.sqrt(event_shape[0].value)) # Image width
nb = kw**2 # Number of basis
nb_half = nb // 2 + 1
nch = 1 # Number of channels in the image
stream = seed_stream.SeedStream(seed=seed, salt="Wishart")
shape = tf.concat([batch_shape, [iw, iw, nb_half - 1]], 0)
# Random sample for the off diagonal values as a dense tensor
x_right = tf.random_normal(shape=shape,
dtype=self.dtype,
seed=stream())
# The upper triangular values needed to get a square kernel per pixel
x_left = tf.zeros(shape)
# Random sample for the diagonal of the matrix
x_diag = tf.random_gamma(shape=[n],
alpha=self._multi_gamma_sequence(
0.5 * self.df, self.p),
beta=0.5,
dtype=self.dtype,
seed=stream())
# Concatenate the diagonal and off-diagonal elements
x_diag = tf.reshape(x_diag, (-1, iw, iw, nch))
x = tf.concat([tf.sqrt(x_diag), x_right], axis=3)
# Scale the sampled matrix using the distribution Scale matrix
diag_scale = tf.exp(self.log_diag_scale)
diag_scale = tf.reshape(diag_scale, (-1, iw, iw, nch))
x *= tf.sqrt(diag_scale) # Square root is equivalent to Cholesky
# Concatenate with the zeros
x = tf.concat([x_left, x], axis=3)
# Create identity basis so that the sampled matrix is only defined by x, if this were not the case
# we would have to do some optimization to find the basis and weights that reconstruct x
identity_basis = tf.eye(num_rows=nb)
identity_basis = tf.reshape(identity_basis, (nb, kw, kw, nch, nch))
sample_shape = tf.concat([batch_shape, [iw, iw, nch]], axis=0)
x_sparse = cov_rep.PrecisionConvCholFilters(
weights_precision=x,
filters_precision=identity_basis,
sample_shape=sample_shape)
return x_sparse
def _sample_n(self, n, seed=None):
seed = seed_stream.SeedStream(seed, "gamma_gamma")
rate = tf.random_gamma(shape=[n],
alpha=self.mixing_concentration,
beta=self.mixing_rate,
dtype=self.dtype,
seed=seed())
return tf.random_gamma(shape=[],
alpha=self.concentration,
beta=rate,
dtype=self.dtype,
seed=seed())
def _sample_n(self, n, seed=None):
seed = seed_stream.SeedStream(seed, "beta")
concentration1 = tf.convert_to_tensor(self.concentration1)
concentration0 = tf.convert_to_tensor(self.concentration0)
shape = self._batch_shape_tensor(concentration1, concentration0)
expanded_concentration1 = tf.broadcast_to(concentration1, shape)
expanded_concentration0 = tf.broadcast_to(concentration0, shape)
gamma1_sample = tf.random.gamma(
shape=[n], alpha=expanded_concentration1, dtype=self.dtype, seed=seed())
gamma2_sample = tf.random.gamma(
shape=[n], alpha=expanded_concentration0, dtype=self.dtype, seed=seed())
beta_sample = gamma1_sample / (gamma1_sample + gamma2_sample)
return beta_sample
def _sample_n(self, n, seed=None):
# Here we use the fact that if:
# lam ~ Gamma(concentration=total_count, rate=(1-probs)/probs)
# then X ~ Poisson(lam) is Negative Binomially distributed.
stream = seed_stream.SeedStream(seed, salt="NegativeBinomial")
rate = tf.random.gamma(
shape=[n],
alpha=self.total_count,
beta=tf.exp(-self.logits),
dtype=self.dtype,
seed=stream())
return tf.random.poisson(
lam=rate, shape=[], dtype=self.dtype, seed=stream())
def _sample_n(self, n, seed=None):
seed = seed_stream.SeedStream(seed, "dirichlet_multinomial")
n_draws = tf.cast(self.total_count, dtype=tf.int32)
k = self.event_shape_tensor()[0]
unnormalized_logits = tf.math.log(
tf.random.gamma(shape=[n],
alpha=self._broadcasted_concentration,
dtype=self.dtype,
seed=seed()))
x = multinomial.draw_sample(1, k, unnormalized_logits, n_draws,
self.dtype, seed())
final_shape = tf.concat([[n], self.batch_shape_tensor(), [k]], 0)
return tf.reshape(x, final_shape)
def _sample_n(self, n, seed=None):
seed = seed_stream.SeedStream(seed, "gamma_gamma")
rate = tf.random.gamma(
shape=[n],
# Be sure to draw enough rates for the fully-broadcasted gamma-gamma.
alpha=self.mixing_concentration +
tf.zeros_like(self.concentration),
beta=self.mixing_rate,
dtype=self.dtype,
seed=seed())
return tf.random.gamma(shape=[],
alpha=self.concentration,
beta=rate,
dtype=self.dtype,
seed=seed())
def sample_distributions(self, sample_shape=(), seed=None, value=None,
name='sample_distributions'):
"""Generate samples and the (random) distributions.
Note that a call to `sample()` without arguments will generate a single
sample.
Args:
sample_shape: 0D or 1D `int32` `Tensor`. Shape of the generated samples.
seed: Python integer seed for generating random numbers.
value: `list` of `Tensor`s in `distribution_fn` order to use to
parameterize other ("downstream") distribution makers.
Default value: `None` (i.e., draw a sample from each distribution).
name: name prepended to ops created by this function.
Default value: `"sample_distributions"`.
Returns:
distributions: a `tuple` of `Distribution` instances for each of
`distribution_fn`.
samples: a `tuple` of `Tensor`s with prepended dimensions `sample_shape`
for each of `distribution_fn`.
"""
seed = seed_stream.SeedStream('JointDistributionSequential', seed)
with self._name_scope(name, values=[sample_shape, seed, value]):
ds = []
if value is None:
xs = [None] * len(self._wrapped_distribution_fn)
else:
xs = list(value)
xs.extend([None]*(len(self._wrapped_distribution_fn) - len(xs)))
if len(xs) != len(self.distribution_fn):
raise ValueError('Number of `xs`s must match number of '
'distributions.')
for i, (dist_fn, args) in enumerate(zip(self._wrapped_distribution_fn,
self._args)):
ds.append(dist_fn(*xs[:i])) # Chain rule of probability.
if xs[i] is None:
# TODO(b/129364796): We should ignore args prefixed with `_`; this
# would mean we more often identify when to use `sample_shape=()`
# rather than `sample_shape=sample_shape`.
xs[i] = ds[-1].sample(
() if args and not self._always_use_specified_sample_shape
else sample_shape, seed=seed())
else:
xs[i] = tf.convert_to_tensor(value=xs[i], dtype_hint=ds[-1].dtype)
seed() # Ensure reproducibility even when xs are (partially) set.
self._most_recently_built_distributions = ds
return tuple(ds), tuple(xs)
def _sample_n(self, n, seed=None):
stream = seed_stream.SeedStream(seed=seed, salt="Wishart")
# Sample a normal full matrix
x = self.normal_dist.sample(sample_shape=n, seed=stream())
# Sample the log diagonal
log_g = self.sqrt_gamma_dist.sample(sample_shape=n, seed=stream())
# Discard the upper triangular part
x = tf.matrix_band_part(x, -1, 0)
# Set the diagonal
x = tf.matrix_set_diag(x, log_g)
return x
def _sample_n(self, n, seed=None):
seed = seed_stream.SeedStream(seed, "beta")
expanded_concentration1 = tf.ones_like(
self.total_concentration, dtype=self.dtype) * self.concentration1
expanded_concentration0 = tf.ones_like(
self.total_concentration, dtype=self.dtype) * self.concentration0
gamma1_sample = tf.random.gamma(shape=[n],
alpha=expanded_concentration1,
dtype=self.dtype,
seed=seed())
gamma2_sample = tf.random.gamma(shape=[n],
alpha=expanded_concentration0,
dtype=self.dtype,
seed=seed())
beta_sample = gamma1_sample / (gamma1_sample + gamma2_sample)
return beta_sample
def _sample_n(self, n, seed=None):
# Like with the univariate Student's t, sampling can be implemented as a
# ratio of samples from a multivariate gaussian with the appropriate
# covariance matrix and a sample from the chi-squared distribution.
seed = seed_stream.SeedStream(seed, salt="multivariate t")
loc = tf.broadcast_to(self.loc, self._sample_shape())
mvn = mvn_linear_operator.MultivariateNormalLinearOperator(
loc=tf.zeros_like(loc), scale=self.scale)
normal_samp = mvn.sample(n, seed=seed())
df = tf.broadcast_to(self.df, self.batch_shape_tensor())
chi2 = chi2_lib.Chi2(df=df)
chi2_samp = chi2.sample(n, seed=seed())
return (self._loc +
normal_samp * tf.math.rsqrt(chi2_samp / self._df)[..., tf.newaxis])
def _sample_n(self, n, seed):
with tf.control_dependencies(self._runtime_assertions):
seed = seed_stream.SeedStream(seed, salt="MixtureSameFamily")
x = self.components_distribution.sample(n, seed=seed()) # [n, B, k, E]
# TODO(jvdillon): Consider using tf.gather (by way of index unrolling).
npdt = x.dtype.as_numpy_dtype
mask = tf.one_hot(
indices=self.mixture_distribution.sample(n, seed=seed()), # [n, B]
depth=self._num_components, # == k
on_value=np.ones([], dtype=npdt),
off_value=np.zeros([], dtype=npdt)) # [n, B, k]
mask = distribution_utils.pad_mixture_dimensions(
mask, self, self.mixture_distribution,
self._event_ndims) # [n, B, k, [1]*e]
x = tf.reduce_sum(x * mask, axis=-1 - self._event_ndims) # [n, B, E]
if self._reparameterize:
x = self._reparameterize_sample(x)
return x
def _sample_n(self, n, seed=None):
# The sampling method comes from the fact that if:
# X ~ Normal(0, 1)
# Z ~ Chi2(df)
# Y = X / sqrt(Z / df)
# then:
# Y ~ StudentT(df).
seed = seed_stream.SeedStream(seed, "student_t")
shape = tf.concat([[n], self.batch_shape_tensor()], 0)
normal_sample = tf.random.normal(shape, dtype=self.dtype, seed=seed())
df = self.df * tf.ones(self.batch_shape_tensor(), dtype=self.dtype)
gamma_sample = tf.random.gamma([n],
0.5 * df,
beta=0.5,
dtype=self.dtype,
seed=seed())
samples = normal_sample * tf.math.rsqrt(gamma_sample / df)
return samples * self.scale + self.loc # Abs(scale) not wanted.
def _sample_n(self, n, seed=None):
seed = seed_stream.SeedStream(seed, "normal_gamma")
shape = tf.concat([[n], self.batch_shape_tensor()], 0)
precision = tf.random_gamma(
shape=shape,
alpha=self.concentration,
beta=self.rate,
dtype=self.dtype,
seed=seed())
scale = tf.sqrt(1 / (self._lambda * precision))
mean = tf.random_normal(
shape=shape, mean=0., stddev=1., dtype=self.loc.dtype, seed=seed())
mean = mean * scale + self.loc
return tf.concat((tf.expand_dims(mean, axis=-1),
tf.expand_dims(precision, axis=-1)), axis=-1)
def _sample_n(self, n, seed=None):
# Get ids as a [n, batch_size]-shaped matrix, unless batch_shape=[] then get
# ids as a [n]-shaped vector.
distributions = self.poisson_and_mixture_distributions()
dist, mixture_dist = distributions
batch_size = tensorshape_util.num_elements(self.batch_shape)
if batch_size is None:
batch_size = tf.reduce_prod(
self._batch_shape_tensor(distributions=distributions))
# We need to 'sample extra' from the mixture distribution if it doesn't
# already specify a probs vector for each batch coordinate.
# We only support this kind of reduced broadcasting, i.e., there is exactly
# one probs vector for all batch dims or one for each.
stream = seed_stream.SeedStream(
seed, salt='PoissonLogNormalQuadratureCompound')
ids = mixture_dist.sample(sample_shape=concat_vectors(
[n],
distribution_util.pick_vector(mixture_dist.is_scalar_batch(),
[batch_size], np.int32([]))),
seed=stream())
# We need to flatten batch dims in case mixture_dist has its own
# batch dims.
ids = tf.reshape(ids,
shape=concat_vectors([n],
distribution_util.pick_vector(
self.is_scalar_batch(),
np.int32([]),
np.int32([-1]))))
# Stride `quadrature_size` for `batch_size` number of times.
offset = tf.range(start=0,
limit=batch_size * self._quadrature_size,
delta=self._quadrature_size,
dtype=ids.dtype)
ids = ids + offset
rate = tf.gather(tf.reshape(dist.rate, shape=[-1]), ids)
rate = tf.reshape(
rate,
shape=concat_vectors(
[n], self._batch_shape_tensor(distributions=distributions)))
return tf.random.poisson(lam=rate,
shape=[],
dtype=self.dtype,
seed=seed)
def _sample_n(self, n, seed=None):
# See
# https://en.wikipedia.org/wiki/Exponentially_modified_Gaussian_distribution
shape = tf.concat([[n], self.batch_shape_tensor()], 0)
stream = seed_stream.SeedStream(seed,
salt="exponentially_modified_normal")
sampled_normal = tf.random_normal(shape=shape,
mean=0.,
stddev=1.,
dtype=self.loc.dtype,
seed=stream())
sampled_uniform = tf.random_uniform(
shape=shape,
minval=np.finfo(self.dtype.as_numpy_dtype).tiny,
maxval=1.,
dtype=self.dtype,
seed=stream())
return (sampled_normal * self.scale + self.loc -
tf.log(sampled_uniform) / self.rate)
def _sample_n(self, n, seed=None):
seed = seed_stream.SeedStream(seed, "dirichlet_multinomial")
n_draws = tf.cast(self.total_count, dtype=tf.int32)
k = self.event_shape_tensor()[0]
unnormalized_logits = tf.reshape(
tf.log(tf.random_gamma(
shape=[n],
alpha=self.concentration,
dtype=self.dtype,
seed=seed())),
shape=[-1, k])
draws = tf.multinomial(
logits=unnormalized_logits,
num_samples=n_draws,
seed=seed())
x = tf.reduce_sum(tf.one_hot(draws, depth=k), -2)
final_shape = tf.concat([[n], self.batch_shape_tensor(), [k]], 0)
x = tf.reshape(x, final_shape)
return tf.cast(x, self.dtype)
def _sample_n(self, n, seed=None):
seed = seed_stream.SeedStream(seed, 'dirichlet_multinomial')
concentration = tf.convert_to_tensor(self._concentration)
total_count = tf.convert_to_tensor(self._total_count)
n_draws = tf.cast(total_count, dtype=tf.int32)
k = self._event_shape_tensor(concentration)[0]
alpha = tf.math.multiply(tf.ones_like(total_count[..., tf.newaxis]),
concentration,
name='alpha')
unnormalized_logits = tf.math.log(
tf.random.gamma(shape=[n],
alpha=alpha,
dtype=self.dtype,
seed=seed()))
x = multinomial.draw_sample(1, k, unnormalized_logits, n_draws,
self.dtype, seed())
final_shape = tf.concat(
[[n],
self._batch_shape_tensor(concentration, total_count), [k]], 0)
return tf.reshape(x, final_shape)
def test_seed_stream(salt='Salt of the Earth', hardcoded_seed=None):
"""Returns a command-line-controllable SeedStream PRNG for unit tests.
When seeding unit-test PRNGs, we want:
- The seed to be fixed to an arbitrary value most of the time, so the test
doesn't flake even if its failure probability is noticeable.
- To switch to different seeds per run when using --runs_per_test to measure
the test's failure probability.
- To set the seed to a specific value when reproducing a low-probability event
(e.g., debugging a crash that only some seeds trigger).
To those ends, this function returns a `SeedStream` seeded with `test_seed`
(which see). The latter respects the command line flags `--fixed_seed=<seed>`
and `--vary-seed` (Boolean, default False). `--vary_seed` uses system entropy
to produce unpredictable seeds. `--fixed_seed` takes precedence over
`--vary_seed` when both are present.
Note that TensorFlow graph mode operations tend to read seed state from two
sources: a "graph-level seed" and an "op-level seed". tf.test.TestCase will
set the former to a fixed value per test, but in general it may be necessary
to explicitly set both to ensure reproducibility.
Args:
salt: Optional string wherewith to salt the returned SeedStream. Setting
this guarantees independent random numbers across tests.
hardcoded_seed: Optional Python value. The seed to use if both the
`--vary_seed` and `--fixed_seed` flags are unset. This should usually be
unnecessary, since a test should pass with any seed.
Returns:
strm: A SeedStream instance seeded with 17, unless otherwise specified by
arguments or command line flags.
"""
return seed_stream.SeedStream(salt, test_seed(hardcoded_seed))
def _flat_sample_distributions(self, sample_shape=(), seed=None, value=None):
"""Executes `model`, creating both samples and distributions."""
ds = []
values_out = []
seed = seed_stream.SeedStream('JointDistributionCoroutine', seed)
gen = self._model()
index = 0
d = next(gen)
if not isinstance(d, self.Root):
raise ValueError('First distribution yielded by coroutine must '
'be wrapped in `Root`.')
try:
while True:
actual_distribution = d.distribution if isinstance(d, self.Root) else d
ds.append(actual_distribution)
if (value is not None and len(value) > index and
value[index] is not None):
seed()
next_value = value[index]
else:
next_value = actual_distribution.sample(
sample_shape=sample_shape if isinstance(d, self.Root) else (),
seed=seed())
if self._validate_args:
with tf.control_dependencies(
self._assert_compatible_shape(
index, sample_shape, next_value)):
values_out.append(tf.identity(next_value))
else:
values_out.append(next_value)
index += 1
d = gen.send(next_value)
except StopIteration:
pass
return ds, values_out
