def sample(self, sample_shape=(), seed=None, name=None):
with tf.name_scope(name or 'sample_particles'):
sample_shape = prefer_static.concat([
dist_util.expand_to_vector(sample_shape), [self.num_particles]
],
axis=0)
return self.distribution.sample(sample_shape, seed=seed)
def sample(self, sample_shape=(), seed=None, name=None):
with tf.name_scope(name or 'sample'):
# Grab the required number of values from the provided tensors.
sample_shape = dist_util.expand_to_vector(sample_shape)
n = tf.cast(tf.reduce_prod(sample_shape), dtype=tf.int32)
# Check that we're not trying to draw too many samples.
assertions = []
will_overflow_ = tf.get_static_value(n > self.max_num_samples)
if will_overflow_:
raise ValueError('Trying to draw {} samples from a '
'`DeterministicEmpirical` instance for which only {} '
'samples were provided.'.format(
tf.get_static_value(n),
tf.get_static_value(self.max_num_samples)))
elif (will_overflow_ is None # Couldn't determine statically.
and self.validate_args):
assertions.append(
tf.debugging.assert_less_equal(
n, self.max_num_samples, message='Number of samples to draw '
'from a `DeterministicEmpirical` instance must not exceed the '
'number provided at construction.'))
# Extract the appropriate number of sampled values.
with tf.control_dependencies(assertions):
sampled = tf.nest.map_structure(
lambda x: x[:n, ...], self.values_with_sample_dim)
# Reshape the values to the appropriate sample shape.
return tf.nest.map_structure(
lambda x: tf.reshape(x, # pylint: disable=g-long-lambda
tf.concat([tf.cast(sample_shape, tf.int32),
tf.cast(tf.shape(x)[1:], tf.int32)],
axis=0)),
sampled)
def new(params, event_shape=(), given_log_count=True,
validate_args=False, name=None):
"""Create the distribution instance from a `params` vector."""
from odin.bay.distributions import ZeroInflated
with tf.compat.v1.name_scope(name, 'ZeroInflatedNegativeBinomial',
[params, event_shape]):
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(
value=event_shape, name='event_shape', dtype=tf.int32),
tensor_name='event_shape')
output_shape = tf.concat([
tf.shape(input=params)[:-1],
event_shape,
], axis=0)
(total_count_params, logits_params,
rate_params) = tf.split(params, 3, axis=-1)
if given_log_count:
total_count_params = tf.exp(total_count_params, name='total_count')
nb = tfd.NegativeBinomial(
total_count=tf.reshape(total_count_params, output_shape),
logits=tf.reshape(logits_params, output_shape),
validate_args=validate_args)
zinb = ZeroInflated(count_distribution=nb,
logits=tf.reshape(rate_params, output_shape),
validate_args=validate_args)
return tfd.Independent(zinb,
reinterpreted_batch_ndims=tf.size(input=event_shape),
validate_args=validate_args)
def new(params,
temperature,
probs_input=False,
event_shape=(),
validate_args=False,
name='RelaxedBernoulliLayer'):
"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(value=event_shape,
name='event_shape',
dtype_hint=tf.int32),
tensor_name='event_shape',
)
new_shape = tf.concat(
[tf.shape(input=params)[:-1], event_shape],
axis=0,
)
params = tf.reshape(params, new_shape)
dist = tfd.Independent(
tfd.RelaxedBernoulli(temperature=temperature,
logits=None if probs_input else params,
probs=params if probs_input else None,
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(input=event_shape),
name=name,
)
return dist
def __init__(
self,
distribution,
sample_shape=(),
validate_args=False,
name=None):
"""Construct the `Sample` distribution.
Args:
distribution: The base distribution instance to transform. Typically an
instance of `Distribution`.
sample_shape: `int` scalar or vector `Tensor` representing the shape of a
single sample.
validate_args: Python `bool`. Whether to validate input with asserts.
If `validate_args` is `False`, and the inputs are invalid,
correct behavior is not guaranteed.
name: The name for ops managed by the distribution.
Default value: `None` (i.e., `'Sample' + distribution.name`).
"""
parameters = dict(locals())
name = name or 'Sample' + distribution.name
self._distribution = distribution
with tf.compat.v1.name_scope(name) as name:
sample_shape = distribution_util.expand_to_vector(tf.convert_to_tensor(
value=sample_shape, dtype_hint=tf.int32, name='sample_shape'))
self._sample_shape = sample_shape
super(Sample, self).__init__(
dtype=self._distribution.dtype,
reparameterization_type=self._distribution.reparameterization_type,
validate_args=validate_args,
allow_nan_stats=self._distribution.allow_nan_stats,
parameters=parameters,
graph_parents=([sample_shape] + distribution._graph_parents), # pylint: disable=protected-access
name=name)
def __init__(self,
distribution,
batch_stack=(),
validate_args=False,
name=None):
"""Construct the `BatchStacker` distribution.
Args:
distribution: The base distribution instance to transform. Typically an
instance of `Distribution`.
batch_stack: `int` scalar or vector `Tensor` representing the shape of a
single sample.
validate_args: Python `bool`. Whether to validate input with asserts.
If `validate_args` is `False`, and the inputs are invalid,
correct behavior is not guaranteed.
name: The name for ops managed by the distribution.
Default value: `None` (i.e., `'BatchStacker' + distribution.name`).
"""
parameters = dict(locals())
name = name or "BatchStacker" + distribution.name
self._distribution = distribution
with tf.name_scope(name) as name:
batch_stack = distribution_util.expand_to_vector(
tf.convert_to_tensor(batch_stack,
dtype_hint=tf.int32,
name="batch_stack"))
self._batch_stack = batch_stack
super(BatchStacker, self).__init__(
dtype=self._distribution.dtype,
reparameterization_type=self._distribution.
reparameterization_type,
validate_args=validate_args,
allow_nan_stats=self._distribution.allow_nan_stats,
parameters=parameters,
name=name,
)
def new(params,
event_shape=(),
dtype=None,
validate_args=False,
name='BernoulliLayer'):
r"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(value=event_shape,
name='event_shape',
dtype_hint=tf.int32),
tensor_name='event_shape',
)
new_shape = tf.concat(
[tf.shape(input=params)[:-1], event_shape],
axis=0,
)
dist = tfd.Independent(
tfd.Bernoulli(logits=tf.reshape(params, new_shape),
dtype=dtype or params.dtype.base_dtype,
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(input=event_shape),
name=name,
)
dist.logits = dist.distribution._logits # pylint: disable=protected-access
dist.probs = dist.distribution._probs # pylint: disable=protected-access
return dist
def new(params,
event_shape=(),
mean_activation=tf.nn.softplus,
disp_activation=softplus1,
validate_args=False,
name="NegativeBinomialDispLayer",
disp=None):
r""" Create the distribution instance from a `params` vector. """
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(value=event_shape,
name='event_shape',
dtype=tf.int32),
tensor_name='event_shape',
)
output_shape = tf.concat(
[tf.shape(input=params)[:-1], event_shape],
axis=0,
)
if disp is None:
loc, disp = tf.split(params, 2, axis=-1)
disp = tf.reshape(disp, output_shape)
else:
loc = params
loc = tf.reshape(loc, output_shape)
loc = mean_activation(loc)
disp = disp_activation(disp)
return tfd.Independent(
NegativeBinomialDisp(loc=loc,
disp=disp,
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(input=event_shape),
name=name,
)
def new(params,
event_shape=(),
given_log_count=True,
validate_args=False,
name=None):
"""Create the distribution instance from a `params` vector."""
from odin.bay.distributions import ZeroInflated
with tf.compat.v1.name_scope(name, 'ZeroInflatedNegativeBinomial',
[params, event_shape]):
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(tf.convert_to_tensor(
value=event_shape, name='event_shape', dtype=tf.int32),
tensor_name='event_shape')
output_shape = tf.concat([
tf.shape(input=params)[:-1],
event_shape,
],
axis=0)
(total_count_params, logits_params,
rate_params) = tf.split(params, 3, axis=-1)
if given_log_count:
total_count_params = tf.exp(total_count_params,
name='total_count')
nb = tfd.NegativeBinomial(
total_count=tf.reshape(total_count_params, output_shape),
logits=tf.reshape(logits_params, output_shape),
validate_args=validate_args)
zinb = ZeroInflated(count_distribution=nb,
logits=tf.reshape(rate_params, output_shape),
validate_args=validate_args)
return tfd.Independent(
zinb,
reinterpreted_batch_ndims=tf.size(input=event_shape),
validate_args=validate_args)
def new(params,
event_shape=(),
alpha_activation=tf.nn.softplus,
beta_activation=tf.nn.softplus,
clip_for_stable=True,
validate_args=False,
name="BetaLayer"):
r"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
alpha_activation = parse_activation(alpha_activation, 'tf')
beta_activation = parse_activation(beta_activation, 'tf')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(value=event_shape,
name='event_shape',
dtype=tf.int32),
tensor_name='event_shape',
)
output_shape = tf.concat(
[tf.shape(input=params)[:-1], event_shape],
axis=0,
)
# alpha, beta
concentration1, concentration0 = tf.split(params, 2, axis=-1)
concentration1 = alpha_activation(concentration1)
concentration0 = beta_activation(concentration0)
if clip_for_stable:
concentration0 = tf.clip_by_value(concentration0, 1e-3, 1e3)
concentration1 = tf.clip_by_value(concentration1, 1e-3, 1e3)
return tfd.Independent(
tfd.Beta(concentration1=tf.reshape(concentration1, output_shape),
concentration0=tf.reshape(concentration0, output_shape),
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(input=event_shape),
name=name,
)
def new(params,
event_shape=(),
softplus_scale=True,
validate_args=False,
name=None):
"""Create the distribution instance from a `params` vector."""
with tf.compat.v1.name_scope(name, 'LogNormal', [params, event_shape]):
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(tf.convert_to_tensor(
value=event_shape, name='event_shape', dtype=tf.int32),
tensor_name='event_shape')
output_shape = tf.concat([
tf.shape(input=params)[:-1],
event_shape,
],
axis=0)
loc_params, scale_params = tf.split(params, 2, axis=-1)
if softplus_scale:
scale_params = tf.math.softplus(
scale_params) + tfd.softplus_inverse(1.0)
return tfd.Independent(
tfd.LogNormal(loc=tf.reshape(loc_params, output_shape),
scale=tf.reshape(scale_params, output_shape),
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(input=event_shape),
validate_args=validate_args)
def new(params, event_shape=(), validate_args=False, name=None):
"""Create the distribution instance from a `params` vector."""
from odin.bay.distributions import ZeroInflated
with tf.compat.v1.name_scope(name, 'ZeroInflatedPoisson',
[params, event_shape]):
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(tf.convert_to_tensor(
value=event_shape, name='event_shape', dtype=tf.int32),
tensor_name='event_shape')
output_shape = tf.concat([
tf.shape(input=params)[:-1],
event_shape,
],
axis=0)
(log_rate_params, logits_params) = tf.split(params, 2, axis=-1)
zip = ZeroInflated(count_distribution=tfd.Poisson(
log_rate=tf.reshape(log_rate_params, output_shape),
validate_args=validate_args),
logits=tf.reshape(logits_params, output_shape),
validate_args=validate_args)
return tfd.Independent(
zip,
reinterpreted_batch_ndims=tf.size(input=event_shape),
validate_args=validate_args)
def new(params,
event_shape=(),
log_prob=None,
reinterpreted_batch_ndims=None,
validate_args=False,
name='DeterministicLayer'):
"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(value=event_shape,
name='event_shape',
dtype=tf.int32),
tensor_name='event_shape',
)
output_shape = tf.concat(
[tf.shape(input=params)[:-1], event_shape],
axis=0,
)
dist = tfd.Deterministic(loc=tf.reshape(params, output_shape),
validate_args=validate_args,
name=name)
# override the log-prob function
if log_prob is not None and callable(log_prob):
dist.log_prob = types.MethodType(log_prob, dist)
# independent
if reinterpreted_batch_ndims is not None and reinterpreted_batch_ndims > 0:
dist = tfd.Independent(
dist, reinterpreted_batch_ndims=int(reinterpreted_batch_ndims))
return dist
def new(params,
event_shape='auto',
pre_softplus=False,
clip_for_stable=True,
validate_args=False,
name=None):
"""Create the distribution instance from a `params` vector."""
event_shape = _preprocess_eventshape(params, event_shape)
with tf.compat.v1.name_scope(name, 'Dirichlet', [params, event_shape]):
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(tf.convert_to_tensor(
value=event_shape, name='event_shape', dtype=tf.int32),
tensor_name='event_shape')
output_shape = tf.concat([
tf.shape(input=params)[:-1],
event_shape,
],
axis=0)
# Clips the Dirichlet parameters to the numerically stable KL region
if pre_softplus:
params = tf.nn.softplus(params)
if clip_for_stable:
params = tf.clip_by_value(params, 1e-3, 1e3)
return tfd.Independent(
tfd.Dirichlet(concentration=tf.reshape(params, output_shape),
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(input=event_shape),
validate_args=validate_args)
def new(params,
event_shape=(),
count_activation=tf.nn.softplus,
validate_args=False,
name='BinomialLayer'):
r"""Create the distribution instance from a `params` vector."""
count_activation = parse_activation(count_activation, 'tf')
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(value=event_shape,
name='event_shape',
dtype=tf.int32),
tensor_name='event_shape',
)
output_shape = tf.concat((tf.shape(params)[:-1], event_shape), axis=0)
total_count, logits = tf.split(params, 2, axis=-1)
total_count = tf.reshape(total_count, output_shape)
logits = tf.reshape(logits, output_shape)
return tfd.Independent(
tfd.Binomial(total_count=count_activation(total_count),
logits=logits,
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(event_shape),
name=name,
)
def new(params,
event_shape=(),
activation=tf.identity,
validate_args=False,
name="ZIPoissonLayer"):
"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(value=event_shape,
name='event_shape',
dtype=tf.int32),
tensor_name='event_shape',
)
output_shape = tf.concat(
[tf.shape(input=params)[:-1], event_shape],
axis=0,
)
(log_rate_params, logits_params) = tf.split(params, 2, axis=-1)
return tfd.Independent(
ZeroInflated(count_distribution=tfd.Poisson(
log_rate=activation(tf.reshape(log_rate_params, output_shape)),
validate_args=validate_args),
logits=tf.reshape(logits_params, output_shape),
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(input=event_shape),
name=name,
)
def new(params,
event_shape,
concentration_activation=softplus1,
rate_activation=softplus1,
validate_args=False,
name="GammaLayer"):
"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(value=event_shape,
name='event_shape',
dtype=tf.int32),
tensor_name='event_shape',
)
output_shape = tf.concat((tf.shape(input=params)[:-1], event_shape), axis=0)
concentration, rate = tf.split(params, 2, axis=-1)
concentration = tf.reshape(concentration, output_shape)
concentration = concentration_activation(concentration)
rate = tf.reshape(rate, output_shape)
rate = rate_activation(rate)
return tfd.Independent(
tfd.Gamma(concentration=concentration,
rate=rate,
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(input=event_shape),
name=name,
)
def new(params,
event_shape=(),
given_logits=True,
validate_args=False,
name='ZIBernoulliLayer'):
"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(value=event_shape,
name='event_shape',
dtype=tf.int32),
tensor_name='event_shape',
)
output_shape = tf.concat(
[tf.shape(input=params)[:-1], event_shape],
axis=0,
)
(bernoulli_params, rate_params) = tf.split(params, 2, axis=-1)
bernoulli_params = tf.reshape(bernoulli_params, output_shape)
bern = tfd.Bernoulli(logits=bernoulli_params if given_logits else None,
probs=bernoulli_params if not given_logits else None,
validate_args=validate_args)
zibern = ZeroInflated(count_distribution=bern,
logits=tf.reshape(rate_params, output_shape),
validate_args=validate_args)
return tfd.Independent(zibern,
reinterpreted_batch_ndims=tf.size(input=event_shape),
name=name)
def new(params,
event_shape=(),
dtype=None,
validate_args=False,
name=None):
"""Create the distribution instance from a `params` vector."""
with tf.name_scope(name, 'IndependentBernoulli',
[params, event_shape]):
params = tf.convert_to_tensor(params, name='params')
event_shape = dist_util.expand_to_vector(tf.convert_to_tensor(
event_shape, name='event_shape', preferred_dtype=tf.int32),
tensor_name='event_shape')
new_shape = tf.concat([
tf.shape(params)[:-1],
event_shape,
],
axis=0)
dist = tfd.Independent(
tfd.Bernoulli(logits=tf.reshape(params, new_shape),
dtype=dtype or params.dtype.base_dtype,
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(event_shape),
validate_args=validate_args)
dist._logits = dist.distribution._logits # pylint: disable=protected-access
dist._probs = dist.distribution._probs # pylint: disable=protected-access
dist.logits = tfd.Bernoulli.logits
dist.probs = tfd.Bernoulli.probs
return dist
def new(params,
event_shape=(),
count_activation=tf.exp,
validate_args=False,
name="NegativeBinomialLayer",
disp=None):
r"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(value=event_shape,
name='event_shape',
dtype=tf.int32),
tensor_name='event_shape',
)
output_shape = tf.concat(
[tf.shape(input=params)[:-1], event_shape],
axis=0,
)
if disp is None:
total_count, logits = tf.split(params, 2, axis=-1)
logits = tf.reshape(logits, output_shape)
else:
total_count = params
logits = disp
total_count = tf.reshape(total_count, output_shape)
total_count = count_activation(total_count)
return tfd.Independent(
tfd.NegativeBinomial(total_count=total_count,
logits=logits,
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(input=event_shape),
name=name,
)
def new(params,
event_shape,
loc_activation,
scale_activation,
validate_args=False,
name="LogNormalLayer"):
"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(value=event_shape,
name='event_shape',
dtype=tf.int32),
tensor_name='event_shape',
)
output_shape = tf.concat(
[tf.shape(input=params)[:-1], event_shape],
axis=0,
)
loc_params, scale_params = tf.split(params, 2, axis=-1)
loc_params = tf.reshape(loc_activation(loc_params), output_shape)
scale_params = tf.reshape(scale_activation(scale_params), output_shape)
return tfd.Independent(
tfd.LogNormal(loc=loc_params,
scale=scale_params,
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(input=event_shape),
name=name,
)
def new(params, event_shape=(), validate_args=False, name=None):
"""Create the distribution instance from a `params` vector."""
from odin.bay.distributions import ZeroInflated
with tf.compat.v1.name_scope(name, 'ZeroInflatedPoisson',
[params, event_shape]):
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(
value=event_shape, name='event_shape', dtype=tf.int32),
tensor_name='event_shape')
output_shape = tf.concat([
tf.shape(input=params)[:-1],
event_shape,
], axis=0)
(log_rate_params, logits_params) = tf.split(params, 2, axis=-1)
zip = ZeroInflated(
count_distribution=tfd.Poisson(
log_rate=tf.reshape(log_rate_params, output_shape),
validate_args=validate_args),
logits=tf.reshape(logits_params, output_shape),
validate_args=validate_args)
return tfd.Independent(zip,
reinterpreted_batch_ndims=tf.size(input=event_shape),
validate_args=validate_args)
def iid_sample(sample_fn, sample_shape):
"""Lift a sampling function to one that draws multiple iid samples.
Args:
sample_fn: Python `callable` that returns a (possibly nested) structure of
`Tensor`s. May optionally take a `seed` named arg: if so, any `int`
seeds (for stateful samplers) are passed through directly, while any
pair-of-`int` seeds (for stateless samplers) are split into independent
seeds for each sample.
sample_shape: `int` `Tensor` shape of iid samples to draw.
Returns:
iid_sample_fn: Python `callable` taking the same arguments as `sample_fn`
and returning iid samples. Each returned `Tensor` will have shape
`concat([sample_shape, shape_of_original_returned_tensor])`.
"""
sample_shape = distribution_util.expand_to_vector(
ps.cast(sample_shape, np.int32), tensor_name='sample_shape')
n = ps.cast(ps.reduce_prod(sample_shape), dtype=np.int32)
def unflatten(x):
unflattened_shape = ps.cast(
ps.concat([sample_shape, ps.shape(x)[1:]], axis=0),
dtype=np.int32)
return tf.reshape(x, unflattened_shape)
def iid_sample_fn(*args, **kwargs):
"""Draws iid samples from `fn`."""
with tf.name_scope('iid_sample_fn'):
seed = kwargs.pop('seed', None)
if samplers.is_stateful_seed(seed):
kwargs = dict(kwargs, seed=SeedStream(seed, salt='iid_sample')())
def pfor_loop_body(_):
with tf.name_scope('iid_sample_fn_stateful_body'):
return sample_fn(*args, **kwargs)
else:
# If a stateless seed arg is passed, split it into `n` different
# stateless seeds, so that we don't just get a bunch of copies of the
# same sample.
if not JAX_MODE:
warnings.warn(
'Saw Tensor seed {}, implying stateless sampling. Autovectorized '
'functions that use stateless sampling may be quite slow because '
'the current implementation falls back to an explicit loop. This '
'will be fixed in the future. For now, you will likely see '
'better performance from stateful sampling, which you can invoke '
'by passing a Python `int` seed.'.format(seed))
seed = samplers.split_seed(seed, n=n, salt='iid_sample_stateless')
def pfor_loop_body(i):
with tf.name_scope('iid_sample_fn_stateless_body'):
return sample_fn(*args, seed=tf.gather(seed, i), **kwargs)
draws = parallel_for.pfor(pfor_loop_body, n)
return tf.nest.map_structure(unflatten, draws, expand_composites=True)
return iid_sample_fn
def iid_sample(sample_fn, sample_shape):
"""Lift a sampling function to one that draws multiple iid samples.
Args:
sample_fn: Python `callable` that returns a (possibly nested) structure of
`Tensor`s. May optionally take a `seed` named arg: if so, any `int`
seeds (for stateful samplers) are passed through directly, while any
pair-of-`int` seeds (for stateless samplers) are split into independent
seeds for each sample.
sample_shape: `int` `Tensor` shape of iid samples to draw.
Returns:
iid_sample_fn: Python `callable` taking the same arguments as `sample_fn`
and returning iid samples. Each returned `Tensor` will have shape
`concat([sample_shape, shape_of_original_returned_tensor])`.
"""
sample_shape = distribution_util.expand_to_vector(
prefer_static.cast(sample_shape, np.int32), tensor_name='sample_shape')
n = prefer_static.cast(prefer_static.reduce_prod(sample_shape),
dtype=np.int32)
def unflatten(x):
unflattened_shape = prefer_static.cast(prefer_static.concat(
[sample_shape, prefer_static.shape(x)[1:]], axis=0),
dtype=np.int32)
return tf.reshape(x, unflattened_shape)
def iid_sample_fn(*args, **kwargs):
"""Draws iid samples from `fn`."""
pfor_loop_body = lambda _: sample_fn(*args, **kwargs)
seed = kwargs.pop('seed', None)
try: # Assume that `seed` is a valid stateful seed (Python `int`).
kwargs = dict(kwargs, seed=SeedStream(seed, salt='iid_sample')())
pfor_loop_body = lambda _: sample_fn(*args, **kwargs)
except TypeError as e:
# If a stateless seed arg is passed, split it into `n` different stateless
# seeds, so that we don't just get a bunch of copies of the same sample.
if TENSOR_SEED_MSG_PREFIX not in str(e):
raise
warnings.warn(
'Saw non-`int` seed {}, implying stateless sampling. '
'Autovectorized functions that use stateless sampling '
'may be quite slow because the current implementation '
'falls back to an explicit loop. This will be fixed in the '
'future. For now, you will likely see better performance '
'from stateful sampling, which you can invoke by passing a'
'traditional Python `int` seed.'.format(seed))
seed = samplers.split_seed(seed, n=n, salt='iid_sample_stateless')
pfor_loop_body = (
lambda i: sample_fn(*args, seed=tf.gather(seed, i), **kwargs))
draws = parallel_for.pfor(pfor_loop_body, n)
return tf.nest.map_structure(unflatten, draws, expand_composites=True)
return iid_sample_fn
def _expand_sample_shape_to_vector(self, x, name):
"""Helper to `sample` which ensures input is 1D."""
x_static_val = tf.contrib.util.constant_value(x)
if x_static_val is None:
prod = tf.reduce_prod(x)
else:
prod = np.prod(x_static_val, dtype=x.dtype.as_numpy_dtype())
x = util.expand_to_vector(x, tensor_name=name)
return x, prod
def _asvi_surrogate_for_sample(dist,
build_nested_surrogate,
sample_shape=None):
"""Builds the surrogate for a `tfd.Sample`-wrapped distribution."""
dist_sample_shape = distribution_util.expand_to_vector(dist.sample_shape)
nested_surrogate = yield from build_nested_surrogate(
dist=dist.distribution,
sample_shape=(dist_sample_shape if sample_shape is None else ps.concat(
[sample_shape, dist_sample_shape], axis=0)))
return independent.Independent(
nested_surrogate,
reinterpreted_batch_ndims=ps.rank_from_shape(dist_sample_shape))
def __init__(self,
shape,
dtype=None,
activation=None,
initializer='zeros',
regularizer=None,
constraint=None,
**kwargs):
"""Creates the `VariableLayer`.
Arguments:
shape: integer or integer vector specifying the shape of the output of
this layer.
dtype: TensorFlow `dtype` of the variable created by this layer.
Default value: `None` (i.e., `tf.as_dtype(tf.keras.backend.floatx())`).
activation: Activation function to use. If you don't specify anything, no
activation is applied (ie. "linear" activation: `a(x) = x`).
Default value: `None`.
initializer: Initializer for the `constant` vector. For example, to
initialize a trainable (initially standard) `Normal` with
`tf.math.softplus` transformed scale, one might use:
```python
tfp.layers.BlockwiseInitializer([
'zeros',
tf.keras.initializers.Constant(np.log(np.expm1(1.))), # = 0.541325
], sizes=[1, 1])
```
Default value: `'zeros'`.
regularizer: Regularizer function applied to the `constant` vector.
Default value: `None`.
constraint: Constraint function applied to the `constant` vector.
Default value: `None`.
**kwargs: Extra arguments forwarded to `tf.keras.layers.Layer`.
"""
super(VariableLayer, self).__init__(**kwargs)
self.activation = tf.keras.activations.get(activation)
self.initializer = tf.keras.initializers.get(initializer)
self.regularizer = tf.keras.regularizers.get(regularizer)
self.constraint = tf.keras.constraints.get(constraint)
shape = tf.get_static_value(dist_util.expand_to_vector(shape))
if shape is None:
raise ValueError('Shape must be known statically.')
self._var = self.add_weight('constant',
shape=shape,
initializer=self.initializer,
regularizer=self.regularizer,
constraint=self.constraint,
dtype=dtype,
trainable=kwargs.get('trainable', True))
def sample(self, sample_shape=(), seed=None, name=None):
with tf.name_scope(name or 'sample_particles'):
sample_shape = prefer_static.concat([
[self.num_particles],
dist_util.expand_to_vector(sample_shape)], axis=0)
x = self.distribution.sample(sample_shape, seed=seed)
def move_particles_to_rightmost_batch_dim(x, event_shape):
ndims = prefer_static.rank_from_shape(prefer_static.shape(x))
event_ndims = prefer_static.rank_from_shape(event_shape)
return dist_util.move_dimension(x, 0, ndims - event_ndims - 1)
return tf.nest.map_structure(
move_particles_to_rightmost_batch_dim,
x, self.distribution.event_shape_tensor())
def new(params,
event_shape=(),
mean_activation=tf.nn.softplus,
disp_activation=softplus1,
validate_args=False,
name="ZINegativeBinomialDispLayer",
disp=None,
rate=None):
r"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(value=event_shape,
name='event_shape',
dtype=tf.int32),
tensor_name='event_shape',
)
output_shape = tf.concat((tf.shape(input=params)[:-1], event_shape),
axis=0)
### splitting the parameters
if disp is None: # full dispersion
if rate is None:
loc, disp, rate = tf.split(params, 3, axis=-1)
rate = tf.reshape(rate, output_shape)
else:
loc, disp = tf.split(params, 2, axis=-1)
disp = tf.reshape(disp, output_shape)
else: # share dispersion
if rate is None:
loc, rate = tf.split(params, 2, axis=-1)
rate = tf.reshape(rate, output_shape)
else:
loc = params
# as count value, do exp if necessary
loc = tf.reshape(loc, output_shape)
loc = mean_activation(loc)
disp = disp_activation(disp)
# create the distribution
nb = NegativeBinomialDisp(loc=loc,
disp=disp,
validate_args=validate_args)
zinb = ZeroInflated(count_distribution=nb,
logits=rate,
validate_args=validate_args)
return tfd.Independent(
zinb,
reinterpreted_batch_ndims=tf.size(input=event_shape),
name=name)
def new(params, event_shape=(), validate_args=False, name=None):
"""Create the distribution instance from a `params` vector."""
with tf.name_scope(name, 'IndependentPoisson', [params, event_shape]):
params = tf.convert_to_tensor(params, name='params')
event_shape = dist_util.expand_to_vector(tf.convert_to_tensor(
event_shape, name='event_shape', preferred_dtype=tf.int32),
tensor_name='event_shape')
output_shape = tf.concat([
tf.shape(params)[:-1],
event_shape,
],
axis=0)
return tfd.Independent(
tfd.Poisson(log_rate=tf.reshape(params, output_shape),
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(event_shape),
validate_args=validate_args)
def __init__(
self,
distribution: tfd.Distribution,
batch_stack=Union[int, Tuple[int, ...], tf.Tensor],
validate_args: bool = False,
name: Optional[str] = None,
):
r"""
Construct the `BatchStacker` distribution.
Parameters
----------
distribution : tfd.Distribution
The base distribution instance to transform. Typically an instance of
``tensorflow_probability.distributions.Distribution``.
batch_stack : Union[int, Tuple[int, ...], tf.Tensor]
Shape of the stack of distributions. To be more precise, ``distribution`` has
its underlying ``batch_shape``, what ``batch_stack`` does effectively is to
add independent replications of ``distribution`` as extra ``batch_shape`` axes,
to the left of the original ``batch_shape``. The resulting batch shape is
``tf.TensorShape(batch_stack) + distribution.batch_shape``.
validate_args : bool
Whether to validate input with asserts. If ``validate_args`` is ``False``,
and the inputs are invalid, correct behavior is not guaranteed.
name : Optional[str]
The name for ops managed by the distribution. If ``None``, defaults to
``"BatchStacker" + distribution.name``.
"""
parameters = dict(locals())
name = name or "BatchStacker" + distribution.name
self._distribution = distribution
with tf.name_scope(name) as name:
batch_stack = distribution_util.expand_to_vector(
tf.convert_to_tensor(batch_stack,
dtype_hint=tf.int32,
name="batch_stack"))
self._batch_stack = batch_stack
super(BatchStacker, self).__init__(
dtype=self._distribution.dtype,
reparameterization_type=self._distribution.
reparameterization_type,
validate_args=validate_args,
allow_nan_stats=self._distribution.allow_nan_stats,
parameters=parameters,
name=name,
)
def new(params,
event_shape=(),
count_activation=tf.exp,
validate_args=False,
name="ZINegativeBinomialLayer",
disp=None,
rate=None):
r"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(value=event_shape,
name='event_shape',
dtype=tf.int32),
tensor_name='event_shape',
)
output_shape = tf.concat((tf.shape(input=params)[:-1], event_shape),
axis=0)
if disp is None: # full dispersion
if rate is None:
total_count, logits, rate = tf.split(params, 3, axis=-1)
rate = tf.reshape(rate, output_shape)
else:
total_count, logits = tf.split(params, 2, axis=-1)
logits = tf.reshape(logits, output_shape)
else: # share dispersion
if rate is None:
total_count, rate = tf.split(params, 2, axis=-1)
rate = tf.reshape(rate, output_shape)
else:
total_count = params
logits = disp
total_count = tf.reshape(total_count, output_shape)
total_count = count_activation(total_count)
nb = tfd.NegativeBinomial(total_count=total_count,
logits=logits,
validate_args=validate_args)
zinb = ZeroInflated(count_distribution=nb,
logits=rate,
validate_args=validate_args)
return tfd.Independent(
zinb,
reinterpreted_batch_ndims=tf.size(input=event_shape),
name=name)
def new(params, event_shape=(), validate_args=False, name=None):
"""Create the distribution instance from a `params` vector."""
with tf.compat.v1.name_scope(name, 'Gamma',
[params, event_shape]):
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(
value=event_shape, name='event_shape', dtype=tf.int32),
tensor_name='event_shape')
output_shape = tf.concat([
tf.shape(input=params)[:-1],
event_shape,
], axis=0)
concentration_params, rate_params = tf.split(params, 2, axis=-1)
return tfd.Independent(
tfd.Gamma(
concentration=tf.reshape(concentration_params, output_shape),
rate=tf.reshape(rate_params, output_shape),
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(input=event_shape),
validate_args=validate_args)
def new(params, event_shape=(), softplus_scale=True,
validate_args=False, name=None):
"""Create the distribution instance from a `params` vector."""
with tf.compat.v1.name_scope(name, 'Normal',
[params, event_shape]):
params = tf.convert_to_tensor(value=params, name='params')
event_shape = dist_util.expand_to_vector(
tf.convert_to_tensor(
value=event_shape, name='event_shape', dtype=tf.int32),
tensor_name='event_shape')
output_shape = tf.concat([
tf.shape(input=params)[:-1],
event_shape,
], axis=0)
loc_params, scale_params = tf.split(params, 2, axis=-1)
if softplus_scale:
scale_params = tf.math.softplus(scale_params) + tfd.softplus_inverse(1.0)
return tfd.Independent(
tfd.Normal(
loc=tf.reshape(loc_params, output_shape),
scale=tf.reshape(scale_params, output_shape),
validate_args=validate_args),
reinterpreted_batch_ndims=tf.size(input=event_shape),
validate_args=validate_args)
