def __init__(self,
count_distribution,
inflated_distribution=None,
logits=None,
probs=None,
eps=1e-8,
validate_args=False,
allow_nan_stats=True,
name="ZeroInflated"):
r"""Initialize a zero-inflated distribution.
A `ZeroInflated` is defined by a zero-inflation rate (`inflated_distribution`,
representing the probabilities of excess zeros) and a `Distribution` object
having matching dtype, batch shape, event shape, and continuity
properties (the dist).
Arguments:
count_distribution : A `tfp.distributions.Distribution` instance.
The instance must have `batch_shape` matching the zero-inflation
distribution.
inflated_distribution: `tfp.distributions.Bernoulli`-like instance.
Manages the probability of excess zeros, the zero-inflated rate.
Must have either scalar `batch_shape` or `batch_shape` matching
`count_distribution.batch_shape`.
logits: An N-D `Tensor` representing the log-odds of a excess zeros
A zero-inflation rate, where the probability of excess zeros is
sigmoid(logits).
Only one of `logits` or `probs` should be passed in.
probs: An N-D `Tensor` representing the probability of a zero event.
Each entry in the `Tensor` parameterizes an independent
ZeroInflated distribution.
Only one of `logits` or `probs` should be passed in.
validate_args: Python `bool`, default `False`. If `True`, raise a runtime
error if batch or event ranks are inconsistent between pi and any of
the distributions. This is only checked if the ranks cannot be
determined statically at graph construction time.
allow_nan_stats: Boolean, default `True`. If `False`, raise an
exception if a statistic (e.g. mean/mode/etc...) is undefined for any
batch member. If `True`, batch members with valid parameters leading to
undefined statistics will return NaN for this statistic.
name: A name for this distribution (optional).
References:
Liu, L. & Blei, D.M.. (2017). Zero-Inflated Exponential Family Embeddings.
Proceedings of the 34th International Conference on Machine Learning,
in PMLR 70:2140-2148
"""
parameters = dict(locals())
with tf.compat.v1.name_scope(name) as name:
# main count distribution
if not isinstance(count_distribution, distribution.Distribution):
raise TypeError(
"count_distribution must be a Distribution instance"
" but saw: %s" % count_distribution)
# Zero inflation distribution
if inflated_distribution is None:
inflated_distribution = Bernoulli(
logits=logits,
probs=probs,
dtype=tf.int32,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
name="ZeroInflatedRate")
elif not isinstance(inflated_distribution,
distribution.Distribution):
raise TypeError(
"inflated_distribution must be a Distribution instance"
" but saw: %s" % inflated_distribution)
# Matching the event shape
inflated_ndim = len(inflated_distribution.event_shape)
count_ndim = len(count_distribution.event_shape)
if inflated_ndim < count_ndim:
inflated_distribution = Independent(inflated_distribution,
count_ndim - inflated_ndim)
self._count_distribution = count_distribution
self._inflated_distribution = inflated_distribution
#
if self._count_distribution.batch_shape.ndims is None:
raise ValueError(
"Expected to know rank(batch_shape) from count_distribution"
)
if self._inflated_distribution.batch_shape.ndims is None:
raise ValueError(
"Expected to know rank(batch_shape) from inflated_distribution"
)
self._eps = tensor_util.convert_nonref_to_tensor(
eps, dtype_hint=count_distribution.dtype, name='eps')
# We let the zero-inflated distribution access _graph_parents since its arguably
# more like a baseclass.
super(ZeroInflated, self).__init__(
dtype=self._count_distribution.dtype,
reparameterization_type=self._count_distribution.
reparameterization_type,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name,
)
def __init__(self, transform_fn, pretransformed_input, dtype=None,
shape=NONE_SPECIFIED, name=None):
"""Creates the `DeferredTensor` object.
Args:
transform_fn: Python `callable` or `tfp.bijectors.Bijector`-like instance.
When `callable`, should take `pretransformed_input` and
return a `Tensor` (representing by this object).
pretransformed_input: object with `shape`, `dtype` properties (typically a
`tf.Variable`) passed into `transform_fn` when this object is acted upon
in a `Tensor` context, eg, `tf.convert_to_tensor`, `+`, `tf.math.exp`,
etc.
dtype: Equivalent to what would otherwise be
`transform_fn(pretransformed_input).dtype`.
Default value: `None` (i.e.,
`getattr(transform_fn, 'dtype', None) or pretransformed_input.dtype`).
shape: Equivalent to what would otherwise be
`transform_fn(pretransformed_input).shape`.
Default value: `'None'` (i.e.,
`getattr(transform_fn, 'forward_event_shape', lambda x: x)(
pretransformed_input.shape)`).
name: Python `str` representing this object's `name`; used only in graph
mode.
Default value: `None` (i.e.,
`(getattr(transform_fn, 'name', None) or
transform_fn.__name__ + '_' + pretransformed_input.name)`).
Raises:
TypeError: if `transform_fn` is not `callable`.
TypeError: if `pretransformed_input` lacks `dtype` and/or `shape`
properties (and `dtype` and/or `shape` arguments are unspecified).
"""
pretransformed_input = tensor_util.convert_nonref_to_tensor(
pretransformed_input,
name='pretransformed_input')
if dtype is None:
dtype = (getattr(transform_fn, 'dtype', None) or
dtype_util.base_dtype(pretransformed_input.dtype))
try:
dtype = None if dtype is None else tf.dtypes.as_dtype(dtype)
except TypeError:
raise TypeError('Argument `dtype` must be convertible to a '
'`tf.dtypes.DType`; saw "{}" of type "{}".'.format(
repr(dtype), type(dtype)))
if shape == NONE_SPECIFIED:
shape = getattr(transform_fn, 'forward_event_shape', _identity)
shape = shape(pretransformed_input.shape)
try:
shape = tf.TensorShape(shape)
except TypeError:
raise TypeError('Argument `shape` must be convertible to a '
'`tf.TensorShape`; saw "{}".'.format(shape))
name = name or getattr(transform_fn, 'name', None)
if not name:
name = '_'.join([
transform_fn.__name__,
getattr(pretransformed_input, 'name', '')])
name = name_util.strip_invalid_chars(name)
name = name_util.camel_to_lower_snake(name)
name = name_util.get_name_scope_name(name)
name = name_util.strip_invalid_chars(name)
if hasattr(transform_fn, 'forward'):
fwd_name = '"{}"'.format(transform_fn.name)
transform_fn = transform_fn.forward
else:
fwd_name = transform_fn.__name__
if not callable(transform_fn):
raise TypeError('Argument `transform_fn` must be `callable`.')
super(DeferredTensor, self).__init__(name=name)
self._transform_fn = transform_fn
self._pretransformed_input = pretransformed_input
self._dtype = dtype
self._shape = shape
self._name = name
self._fwd_name = fwd_name
# Secret handshake with tf.is_tensor to return True for DT.
#
# Works around an exception in LinearOperator (which in 2.0.0 checks only
# `tf.is_tensor`, not also `linear_operator_util.is_ref`:
# ValueError: Graph parent item 0 is not a Tensor;
# <DeferredTensor: dtype=float32, shape=[2], fn=exp>.
# TODO(b/140157055): Remove this shim after LinOp is patched in 2.0.
self.is_tensor_like = True
def __init__(self,
rate1=None,
rate2=None,
log_rate1=None,
log_rate2=None,
force_probs_to_zero_outside_support=False,
validate_args=False,
allow_nan_stats=True,
name='Skellam'):
"""Initialize a batch of Skellam distributions.
Args:
rate1: Floating point tensor, the first rate parameter. `rate1` must be
positive. Must specify exactly one of `rate1` and `log_rate1`
rate2: Floating point tensor, the second rate parameter. `rate` must be
positive. Must specify exactly one of `rate2` and `log_rate2`.
log_rate1: Floating point tensor, the log of the first rate parameter.
Must specify exactly one of `rate1` and `log_rate1`.
log_rate2: Floating point tensor, the log of the second rate parameter.
Must specify exactly one of `rate2` and `log_rate2`.
force_probs_to_zero_outside_support: Python `bool`. When `True`,
`log_prob` returns `-inf` (and `prob` returns `0`) for non-integer
inputs. When `False`, `log_prob` evaluates the Skellam pmf as a
continuous function (note that this function is not itself
a normalized probability log-density).
Default value: `False`.
validate_args: Python `bool`. 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`. 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.
Raises:
ValueError: if none or both of `rate1`, `log_rate1` are specified.
ValueError: if none or both of `rate2`, `log_rate2` are specified.
"""
parameters = dict(locals())
if (rate1 is None) == (log_rate1 is None):
raise ValueError(
'Must specify exactly one of `rate1` and `log_rate1`.')
if (rate2 is None) == (log_rate2 is None):
raise ValueError(
'Must specify exactly one of `rate2` and `log_rate2`.')
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype(
[rate1, rate2, log_rate1, log_rate2], dtype_hint=tf.float32)
self._rate1 = tensor_util.convert_nonref_to_tensor(rate1,
name='rate1',
dtype=dtype)
self._log_rate1 = tensor_util.convert_nonref_to_tensor(
log_rate1, name='log_rate1', dtype=dtype)
self._rate2 = tensor_util.convert_nonref_to_tensor(rate2,
name='rate2',
dtype=dtype)
self._log_rate2 = tensor_util.convert_nonref_to_tensor(
log_rate2, name='log_rate2', dtype=dtype)
self._force_probs_to_zero_outside_support = force_probs_to_zero_outside_support
super(Skellam, self).__init__(
dtype=dtype,
reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name)
def test_np_object(self):
x = np.array(0.)
y = tensor_util.convert_nonref_to_tensor(x)
self.assertIsInstance(y, tf.Tensor)
self.assertEqual(x, self.evaluate(y))
def __init__(self,
loc,
scale_row,
scale_column,
validate_args=False,
allow_nan_stats=True,
name='MatrixNormalLinearOperator'):
"""Construct Matrix Normal distribution on `R^{n x p}`.
The `batch_shape` is the broadcast shape between `loc`, `scale_row`
and `scale_column` arguments.
The `event_shape` is given by the matrix implied by `loc`.
Args:
loc: Floating-point `Tensor`, having shape `[B1, ..., Bb, n, p]`.
scale_row: Instance of `LinearOperator` with the same `dtype` as `loc`
and shape `[B1, ..., Bb, n, n]`.
scale_column: Instance of `LinearOperator` with the same `dtype` as `loc`
and shape `[B1, ..., Bb, p, p]`.
validate_args: Python `bool`, default `False`. Whether to validate input
with asserts. If `validate_args` is `False`, and the inputs are
invalid, correct behavior is not guaranteed.
allow_nan_stats: Python `bool`, default `True`. If `False`, raise an
exception if a statistic (e.g. mean/mode/etc...) is undefined for any
batch member If `True`, batch members with valid parameters leading to
undefined statistics will return NaN for this statistic.
name: The name to give Ops created by the initializer.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([loc, scale_column, scale_row],
dtype_hint=tf.float32)
loc = tensor_util.convert_nonref_to_tensor(loc,
dtype=dtype,
name='loc')
self._loc = loc
if not hasattr(scale_row, 'matmul'):
raise ValueError(
'`scale_row` must be a `tf.linalg.LinearOperator`.')
if not hasattr(scale_column, 'matmul'):
raise ValueError(
'`scale_column` must be a `tf.linalg.LinearOperator`.')
if validate_args and not scale_row.is_non_singular:
raise ValueError('`scale_row` must be non-singular.')
if validate_args and not scale_column.is_non_singular:
raise ValueError('`scale_column` must be non-singular.')
self._scale_row = scale_row
self._scale_column = scale_column
super(MatrixNormalLinearOperator, self).__init__(
dtype=dtype,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
parameters=parameters,
name=name)
self._parameters = parameters
def __init__(self,
loc,
scale,
low,
high,
validate_args=False,
allow_nan_stats=True,
name='TruncatedNormal'):
"""Construct TruncatedNormal.
All parameters of the distribution will be broadcast to the same shape,
so the resulting distribution will have a batch_shape of the broadcast
shape of all parameters.
Args:
loc: Floating point tensor; the mean of the normal distribution(s) (
note that the mean of the resulting distribution will be different
since it is modified by the bounds).
scale: Floating point tensor; the std deviation of the normal
distribution(s).
low: `float` `Tensor` representing lower bound of the distribution's
support. Must be such that `low < high`.
high: `float` `Tensor` representing upper bound of the distribution's
support. Must be such that `low < high`.
validate_args: Python `bool`, default `False`. When `True` distribution
parameters are checked at run-time.
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([loc, scale, low, high],
tf.float32)
self._loc = tensor_util.convert_nonref_to_tensor(loc,
name='loc',
dtype=dtype)
self._scale = tensor_util.convert_nonref_to_tensor(scale,
name='scale',
dtype=dtype)
self._low = tensor_util.convert_nonref_to_tensor(low,
name='low',
dtype=dtype)
self._high = tensor_util.convert_nonref_to_tensor(high,
name='high',
dtype=dtype)
dtype_util.assert_same_float_dtype(
[self._loc, self._scale, self._low, self._high])
super(TruncatedNormal, self).__init__(
dtype=dtype,
# This distribution is fully reparameterized. loc, scale have straight
# through gradients. The gradients for the bounds are implemented
# using custom derived expressions based on implicit gradients.
# For the special case of lower bound zero and a positive upper bound
# an equivalent expression can also be found in Sec 9.1.1.
# of https://arxiv.org/pdf/1806.01851.pdf. The implementation here
# handles arbitrary bounds.
reparameterization_type=reparameterization.
FULLY_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name)
def __init__(self,
low=None,
high=None,
hinge_softness=None,
validate_args=False,
name='soft_clip'):
"""Instantiates the SoftClip bijector.
Args:
low: Optional float `Tensor` lower bound. If `None`, the lower-bound
constraint is omitted.
Default value: `None`.
high: Optional float `Tensor` upper bound. If `None`, the upper-bound
constraint is omitted.
Default value: `None`.
hinge_softness: Optional nonzero float `Tensor`. Controls the softness
of the constraint at the boundaries; values outside of the constraint
set are mapped into intervals of width approximately
`log(2) * hinge_softness` on the interior of each boundary. High
softness reserves more space for values outside of the constraint set,
leading to greater distortion of inputs *within* the constraint set,
but improved numerical stability near the boundaries.
Default value: `None` (`1.0`).
validate_args: Python `bool` indicating whether arguments should be
checked for correctness.
name: Python `str` name given to ops managed by this object.
"""
with tf.name_scope(name):
dtype = dtype_util.common_dtype([low, high, hinge_softness],
dtype_hint=tf.float32)
low = tensor_util.convert_nonref_to_tensor(low,
name='low',
dtype=dtype)
high = tensor_util.convert_nonref_to_tensor(high,
name='high',
dtype=dtype)
hinge_softness = tensor_util.convert_nonref_to_tensor(
hinge_softness, name='hinge_softness', dtype=dtype)
softplus_bijector = softplus.Softplus(
hinge_softness=hinge_softness)
negate = tf.convert_to_tensor(-1., dtype=dtype)
components = []
if low is not None and high is not None:
# Support reference tensors (eg Variables) for `high` and `low` by
# deferring all computation on them until needed.
width = tfp_util.DeferredTensor(
pretransformed_input=high,
transform_fn=lambda high: high - low)
negated_shrinkage_factor = tfp_util.DeferredTensor(
pretransformed_input=width,
transform_fn=lambda w: tf.cast( # pylint: disable=g-long-lambda
negate * w / softplus_bijector.forward(w),
dtype=dtype))
# Implement the soft constraint from 'Mathematical Details' above:
# softclip(x) := -softplus(width - softplus(x - low)) *
# (width) / (softplus(width)) + high
components = [
shift.Shift(high),
scale.Scale(negated_shrinkage_factor), softplus_bijector,
shift.Shift(width),
scale.Scale(negate), softplus_bijector,
shift.Shift(tfp_util.DeferredTensor(low, lambda x: -x))
]
elif low is not None:
# Implement a soft lower bound:
# softlower(x) := softplus(x - low) + low
components = [
shift.Shift(low), softplus_bijector,
shift.Shift(tfp_util.DeferredTensor(low, lambda x: -x))
]
elif high is not None:
# Implement a soft upper bound:
# softupper(x) := -softplus(high - x) + high
components = [
shift.Shift(high),
scale.Scale(negate), softplus_bijector,
scale.Scale(negate),
shift.Shift(high)
]
self._low = low
self._high = high
self._hinge_softness = hinge_softness
self._chain = chain.Chain(components, validate_args=validate_args)
super(SoftClip, self).__init__(forward_min_event_ndims=0,
dtype=dtype,
validate_args=validate_args,
is_constant_jacobian=not components,
name=name)
def __init__(self,
mean_direction,
concentration,
validate_args=False,
allow_nan_stats=True,
name='VonMisesFisher'):
"""Creates a new `VonMisesFisher` instance.
Args:
mean_direction: Floating-point `Tensor` with shape [B1, ... Bn, D].
A unit vector indicating the mode of the distribution, or the
unit-normalized direction of the mean. (This is *not* in general the
mean of the distribution; the mean is not generally in the support of
the distribution.) NOTE: `D` is currently restricted to <= 5.
concentration: Floating-point `Tensor` having batch shape [B1, ... Bn]
broadcastable with `mean_direction`. The level of concentration of
samples around the `mean_direction`. `concentration=0` indicates a
uniform distribution over the unit hypersphere, and `concentration=+inf`
indicates a `Deterministic` distribution (delta function) at
`mean_direction`.
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.
Raises:
ValueError: For known-bad arguments, i.e. unsupported event dimension.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([mean_direction, concentration],
tf.float32)
self._mean_direction = tensor_util.convert_nonref_to_tensor(
mean_direction, name='mean_direction', dtype=dtype)
self._concentration = tensor_util.convert_nonref_to_tensor(
concentration, name='concentration', dtype=dtype)
static_event_dim = tf.compat.dimension_value(
tensorshape_util.with_rank_at_least(self._mean_direction.shape,
1)[-1])
if static_event_dim is not None and static_event_dim > 5:
raise ValueError('von Mises-Fisher ndims > 5 is not currently '
'supported')
# mean_direction is always reparameterized.
# concentration is only for event_dim==3, via an inversion sampler.
reparameterization_type = (reparameterization.FULLY_REPARAMETERIZED
if static_event_dim == 3 else
reparameterization.NOT_REPARAMETERIZED)
super(VonMisesFisher, self).__init__(
dtype=self._concentration.dtype,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
reparameterization_type=reparameterization_type,
parameters=parameters,
name=name)
def __init__(self,
initial_distribution,
transition_distribution,
observation_distribution,
num_steps,
validate_args=False,
allow_nan_stats=True,
time_varying_observation_distribution=False,
name='HiddenMarkovModel'):
"""Initialize hidden Markov model.
Args:
initial_distribution: A `Categorical`-like instance.
Determines probability of first hidden state in Markov chain.
The number of categories must match the number of categories of
`transition_distribution` as well as both the rightmost batch
dimension of `transition_distribution` and the rightmost batch
dimension of `observation_distribution`.
transition_distribution: A `Categorical`-like instance.
The rightmost batch dimension indexes the probability distribution
of each hidden state conditioned on the previous hidden state.
observation_distribution: A `tfp.distributions.Distribution`-like
instance. The rightmost batch dimension indexes the distribution
of each observation conditioned on the corresponding hidden state.
num_steps: The number of steps taken in Markov chain. An integer valued
tensor. The number of transitions is `num_steps - 1`.
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`.
time_varying_observation_distribution: Python `bool`, default `False`.
When `True`, the observation_distribution has an additional batch
dimension that indexes the distribution of each observation conditioned
on the corresponding timestep. This dimension size should always match
num_steps and is the second-to-last batch axis in the batch dimensions
(just to the left of the dimension for the number of states).
name: Python `str` name prefixed to Ops created by this class.
Default value: "HiddenMarkovModel".
Raises:
ValueError: if `num_steps` is not at least 1.
ValueError: if `initial_distribution` does not have scalar `event_shape`.
ValueError: if `transition_distribution` does not have scalar
`event_shape.`
ValueError: if `transition_distribution` and `observation_distribution`
are fully defined but don't have matching rightmost dimension.
"""
parameters = dict(locals())
# pylint: disable=protected-access
with tf.name_scope(name) as name:
self._num_steps = tensor_util.convert_nonref_to_tensor(num_steps)
self._initial_distribution = initial_distribution
self._observation_distribution = observation_distribution
self._transition_distribution = transition_distribution
self._time_varying_observation_distribution = (
time_varying_observation_distribution)
num_steps_ = tf.get_static_value(num_steps)
if num_steps_ is not None:
if np.ndim(num_steps_) != 0:
raise ValueError(
'`num_steps` must be a scalar but it has rank {}'.format(
np.ndim(num_steps_)))
else:
self._static_event_shape = tf.TensorShape(
[num_steps_]).concatenate(
self._observation_distribution.event_shape)
else:
self._static_event_shape = tf.TensorShape(
[None]).concatenate(
self._observation_distribution.event_shape)
observation_batch_shape = (
self._observation_distribution.batch_shape[:-2]
if self._time_varying_observation_distribution
else self._observation_distribution.batch_shape[:-1])
self._static_batch_shape = tf.broadcast_static_shape(
self._initial_distribution.batch_shape,
tf.broadcast_static_shape(
self._transition_distribution.batch_shape[:-1],
observation_batch_shape))
# pylint: disable=protected-access
super(HiddenMarkovModel, self).__init__(
dtype=self._observation_distribution.dtype,
reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name)
# pylint: enable=protected-access
self._parameters = parameters
def __init__(self,
distribution,
bijector,
batch_shape=None,
event_shape=None,
kwargs_split_fn=_default_kwargs_split_fn,
validate_args=False,
parameters=None,
name=None):
"""Construct a Transformed Distribution.
Args:
distribution: The base distribution instance to transform. Typically an
instance of `Distribution`.
bijector: The object responsible for calculating the transformation.
Typically an instance of `Bijector`.
batch_shape: `integer` vector `Tensor` which overrides `distribution`
`batch_shape`; valid only if `distribution.is_scalar_batch()`.
event_shape: `integer` vector `Tensor` which overrides `distribution`
`event_shape`; valid only if `distribution.is_scalar_event()`.
kwargs_split_fn: Python `callable` which takes a kwargs `dict` and returns
a tuple of kwargs `dict`s for each of the `distribution` and `bijector`
parameters respectively.
Default value: `_default_kwargs_split_fn` (i.e.,
`lambda kwargs: (kwargs.get('distribution_kwargs', {}),
kwargs.get('bijector_kwargs', {}))`)
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.
parameters: Locals dict captured by subclass constructor, to be used for
copy/slice re-instantiation operations.
name: Python `str` name prefixed to Ops created by this class. Default:
`bijector.name + distribution.name`.
Raises:
ValueError: If `distribution` is a joint distribution and a `batch_shape`
override is passed.
ValueError: If `distribution` is a joint distribution and an `event_shape`
override is passed.
"""
parameters = dict(locals()) if parameters is None else parameters
name = name or (('' if bijector is None else bijector.name) +
(distribution.name or ''))
with tf.name_scope(name) as name:
self._distribution = distribution
self._bijector = bijector
self._kwargs_split_fn = (_default_kwargs_split_fn
if kwargs_split_fn is None else
kwargs_split_fn)
# For convenience we define some handy constants.
self._zero = tf.constant(0, dtype=tf.int32, name='zero')
self._empty = tf.constant([], dtype=tf.int32, name='empty')
# We don't just want to check isinstance(JointDistribution) because
# TransformedDistributions with multipart bijectors are effectively
# joint but don't inherit from JD. The 'duck-type' test is that
# JDs have a structured dtype.
self._base_is_joint = tf.nest.is_nested(self.distribution.dtype)
if self._base_is_joint:
if batch_shape:
raise ValueError(
'Overriding the batch shape of a joint distribution'
' ({}) is not currently supported.'.format(
self.distribution))
if event_shape:
raise ValueError(
'Overriding the event shape of a joint distribution'
' ({}) is not currently supported.'.format(
self.distribution))
override_batch_shape = self._empty if batch_shape is None else batch_shape
self._override_batch_shape = tensor_util.convert_nonref_to_tensor(
override_batch_shape,
dtype=tf.int32,
name='override_batch_shape')
override_event_shape = self._empty if event_shape is None else event_shape
self._override_event_shape = tensor_util.convert_nonref_to_tensor(
override_event_shape,
dtype=tf.int32,
name='override_event_shape')
# `_is_maybe_{batch, event}_override` is False if we know statically that
# the batch/event shape is not being overridden; otherwise it is True.
self._is_maybe_event_override = tensorshape_util.dims(
self._override_event_shape.shape) != [0]
self._is_maybe_batch_override = tensorshape_util.dims(
self._override_batch_shape.shape) != [0]
dtype = self.bijector.forward_dtype(self.distribution.dtype)
self._is_joint = tf.nest.is_nested(dtype)
super(TransformedDistribution, self).__init__(
dtype=dtype,
reparameterization_type=self._distribution.
reparameterization_type,
validate_args=validate_args,
allow_nan_stats=self._distribution.allow_nan_stats,
parameters=parameters,
name=name)
def __init__(self,
kernel,
index_points=None,
observation_index_points=None,
observations=None,
observation_noise_variance=0.,
predictive_noise_variance=None,
mean_fn=None,
cholesky_fn=None,
jitter=1e-6,
validate_args=False,
allow_nan_stats=False,
name='GaussianProcessRegressionModel',
_conditional_kernel=None,
_conditional_mean_fn=None):
"""Construct a GaussianProcessRegressionModel instance.
Args:
kernel: `PositiveSemidefiniteKernel`-like instance representing the
GP's covariance function.
index_points: `float` `Tensor` representing finite collection, or batch of
collections, of points in the index set over which the GP is defined.
Shape has the form `[b1, ..., bB, e, f1, ..., fF]` where `F` is the
number of feature dimensions and must equal `kernel.feature_ndims` and
`e` is the number (size) of index points in each batch. Ultimately this
distribution corresponds to an `e`-dimensional multivariate normal. The
batch shape must be broadcastable with `kernel.batch_shape` and any
batch dims yielded by `mean_fn`.
observation_index_points: `float` `Tensor` representing finite collection,
or batch of collections, of points in the index set for which some data
has been observed. Shape has the form `[b1, ..., bB, e, f1, ..., fF]`
where `F` is the number of feature dimensions and must equal
`kernel.feature_ndims`, and `e` is the number (size) of index points in
each batch. `[b1, ..., bB, e]` must be broadcastable with the shape of
`observations`, and `[b1, ..., bB]` must be broadcastable with the
shapes of all other batched parameters (`kernel.batch_shape`,
`index_points`, etc). The default value is `None`, which corresponds to
the empty set of observations, and simply results in the prior
predictive model (a GP with noise of variance
`predictive_noise_variance`).
observations: `float` `Tensor` representing collection, or batch of
collections, of observations corresponding to
`observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
must be brodcastable with the batch and example shapes of
`observation_index_points`. The batch shape `[b1, ..., bB]` must be
broadcastable with the shapes of all other batched parameters
(`kernel.batch_shape`, `index_points`, etc.). The default value is
`None`, which corresponds to the empty set of observations, and simply
results in the prior predictive model (a GP with noise of variance
`predictive_noise_variance`).
observation_noise_variance: `float` `Tensor` representing the variance
of the noise in the Normal likelihood distribution of the model. May be
batched, in which case the batch shape must be broadcastable with the
shapes of all other batched parameters (`kernel.batch_shape`,
`index_points`, etc.).
Default value: `0.`
predictive_noise_variance: `float` `Tensor` representing the variance in
the posterior predictive model. If `None`, we simply re-use
`observation_noise_variance` for the posterior predictive noise. If set
explicitly, however, we use this value. This allows us, for example, to
omit predictive noise variance (by setting this to zero) to obtain
noiseless posterior predictions of function values, conditioned on noisy
observations.
mean_fn: Python `callable` that acts on `index_points` to produce a
collection, or batch of collections, of mean values at `index_points`.
Takes a `Tensor` of shape `[b1, ..., bB, f1, ..., fF]` and returns a
`Tensor` whose shape is broadcastable with `[b1, ..., bB]`.
Default value: `None` implies the constant zero function.
cholesky_fn: Callable which takes a single (batch) matrix argument and
returns a Cholesky-like lower triangular factor. Default value: `None`,
in which case `make_cholesky_with_jitter_fn` is used with the `jitter`
parameter.
jitter: `float` scalar `Tensor` added to the diagonal of the covariance
matrix to ensure positive definiteness of the covariance matrix.
This argument is ignored if `cholesky_fn` is set.
Default value: `1e-6`.
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: `False`.
name: Python `str` name prefixed to Ops created by this class.
Default value: 'GaussianProcessRegressionModel'.
_conditional_kernel: Internal parameter -- do not use.
_conditional_mean_fn: Internal parameter -- do not use.
Raises:
ValueError: if either
- only one of `observations` and `observation_index_points` is given, or
- `mean_fn` is not `None` and not callable.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([
index_points, observation_index_points, observations,
observation_noise_variance, predictive_noise_variance, jitter
], tf.float32)
index_points = tensor_util.convert_nonref_to_tensor(
index_points, dtype=dtype, name='index_points')
observation_index_points = tensor_util.convert_nonref_to_tensor(
observation_index_points,
dtype=dtype,
name='observation_index_points')
observations = tensor_util.convert_nonref_to_tensor(
observations, dtype=dtype, name='observations')
observation_noise_variance = tensor_util.convert_nonref_to_tensor(
observation_noise_variance,
dtype=dtype,
name='observation_noise_variance')
predictive_noise_variance = tensor_util.convert_nonref_to_tensor(
predictive_noise_variance,
dtype=dtype,
name='predictive_noise_variance')
if predictive_noise_variance is None:
predictive_noise_variance = observation_noise_variance
jitter = tensor_util.convert_nonref_to_tensor(jitter,
dtype=dtype,
name='jitter')
if (observation_index_points is None) != (observations is None):
raise ValueError(
'`observations` and `observation_index_points` must both be given '
'or None. Got {} and {}, respectively.'.format(
observations, observation_index_points))
# Default to a constant zero function, borrowing the dtype from
# index_points to ensure consistency.
if mean_fn is None:
mean_fn = lambda x: tf.zeros([1], dtype=dtype)
else:
if not callable(mean_fn):
raise ValueError('`mean_fn` must be a Python callable')
if cholesky_fn is None:
cholesky_fn = cholesky_util.make_cholesky_with_jitter_fn(
jitter)
self._name = name
self._observation_index_points = observation_index_points
self._observations = observations
self._observation_noise_variance = observation_noise_variance
self._predictive_noise_variance = predictive_noise_variance
self._jitter = jitter
self._validate_args = validate_args
with tf.name_scope('init'):
if _conditional_kernel is None:
_conditional_kernel = tfpk.SchurComplement(
base_kernel=kernel,
fixed_inputs=observation_index_points,
cholesky_fn=cholesky_fn,
diag_shift=observation_noise_variance)
# Special logic for mean_fn only; SchurComplement already handles the
# case of empty observations (ie, falls back to base_kernel).
if _is_empty_observation_data(
feature_ndims=kernel.feature_ndims,
observation_index_points=observation_index_points,
observations=observations):
if _conditional_mean_fn is None:
_conditional_mean_fn = mean_fn
else:
_validate_observation_data(
kernel=kernel,
observation_index_points=observation_index_points,
observations=observations)
if _conditional_mean_fn is None:
def conditional_mean_fn(x):
"""Conditional mean."""
observations = tf.convert_to_tensor(
self._observations)
observation_index_points = tf.convert_to_tensor(
self._observation_index_points)
k_x_obs_linop = tf.linalg.LinearOperatorFullMatrix(
kernel.matrix(x, observation_index_points))
chol_linop = tf.linalg.LinearOperatorLowerTriangular(
_conditional_kernel.divisor_matrix_cholesky(
fixed_inputs=observation_index_points))
diff = observations - mean_fn(
observation_index_points)
return mean_fn(x) + k_x_obs_linop.matvec(
chol_linop.solvevec(chol_linop.solvevec(diff),
adjoint=True))
_conditional_mean_fn = conditional_mean_fn
super(GaussianProcessRegressionModel, self).__init__(
kernel=_conditional_kernel,
mean_fn=_conditional_mean_fn,
index_points=index_points,
cholesky_fn=cholesky_fn,
jitter=jitter,
# What the GP super class calls "observation noise variance" we call
# here the "predictive noise variance". We use the observation noise
# variance for the fit/solve process above, and predictive for
# downstream computations like sampling.
observation_noise_variance=predictive_noise_variance,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
name=name)
self._parameters = parameters
def __init__(self,
kernel,
index_points=None,
mean_fn=None,
observation_noise_variance=0.,
marginal_fn=None,
cholesky_fn=None,
jitter=1e-6,
always_yield_multivariate_normal=False,
validate_args=False,
allow_nan_stats=False,
parameters=None,
name='GaussianProcess',
_check_marginal_cholesky_fn=True):
"""Instantiate a GaussianProcess Distribution.
Args:
kernel: `PositiveSemidefiniteKernel`-like instance representing the
GP's covariance function.
index_points: `float` `Tensor` representing finite (batch of) vector(s) of
points in the index set over which the GP is defined. Shape has the form
`[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature
dimensions and must equal `kernel.feature_ndims` and `e` is the number
(size) of index points in each batch. Ultimately this distribution
corresponds to a `e`-dimensional multivariate normal. The batch shape
must be broadcastable with `kernel.batch_shape` and any batch dims
yielded by `mean_fn`.
mean_fn: Python `callable` that acts on `index_points` to produce a (batch
of) vector(s) of mean values at `index_points`. Takes a `Tensor` of
shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is
broadcastable with `[b1, ..., bB]`. Default value: `None` implies
constant zero function.
observation_noise_variance: `float` `Tensor` representing (batch of)
scalar variance(s) of the noise in the Normal likelihood
distribution of the model. If batched, the batch shape must be
broadcastable with the shapes of all other batched parameters
(`kernel.batch_shape`, `index_points`, etc.).
Default value: `0.`
marginal_fn: A Python callable that takes a location, covariance matrix,
optional `validate_args`, `allow_nan_stats` and `name` arguments, and
returns a multivariate normal subclass of `tfd.Distribution`.
At most one of `cholesky_fn` and `marginal_fn` should be set.
Default value: `None`, in which case a Cholesky-factorizing function
is created using `make_cholesky_factored_marginal_fn` and the
`cholesky_fn` argument.
cholesky_fn: Callable which takes a single (batch) matrix argument and
returns a Cholesky-like lower triangular factor. Default value: `None`,
in which case `make_cholesky_with_jitter_fn` is used with the `jitter`
parameter. At most one of `cholesky_fn` and `marginal_fn` should be set.
jitter: `float` scalar `Tensor` added to the diagonal of the covariance
matrix to ensure positive definiteness of the covariance matrix, when
`marginal_fn` and `cholesky_fn` is None.
This argument is ignored if `cholesky_fn` is set.
Default value: `1e-6`.
always_yield_multivariate_normal: If `False` (the default), we produce a
scalar `Normal` distribution when the number of `index_points` is
statically known to be `1`. If `True`, we avoid this behavior, ensuring
that the event shape will retain the `1` from `index_points`.
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: `False`.
parameters: For subclasses, a dict of constructor arguments.
name: Python `str` name prefixed to Ops created by this class.
Default value: "GaussianProcess".
_check_marginal_cholesky_fn: Internal parameter -- do not use.
Raises:
ValueError: if `mean_fn` is not `None` and is not callable.
"""
parameters = dict(locals()) if parameters is None else parameters
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype(
{
'index_points': index_points,
'observation_noise_variance': observation_noise_variance,
'jitter': jitter
}, tf.float32)
index_points = tensor_util.convert_nonref_to_tensor(
index_points, dtype=dtype, name='index_points')
jitter = tensor_util.convert_nonref_to_tensor(jitter,
dtype=dtype,
name='jitter')
observation_noise_variance = tensor_util.convert_nonref_to_tensor(
observation_noise_variance,
dtype=dtype,
name='observation_noise_variance')
self._kernel = kernel
self._index_points = index_points
# Default to a constant zero function, borrowing the dtype from
# index_points to ensure consistency.
if mean_fn is None:
mean_fn = lambda x: tf.zeros([1], dtype=dtype)
else:
if not callable(mean_fn):
raise ValueError('`mean_fn` must be a Python callable')
self._mean_fn = mean_fn
self._observation_noise_variance = observation_noise_variance
self._jitter = jitter
self._cholesky_fn = cholesky_fn
if (_check_marginal_cholesky_fn and marginal_fn is not None
and cholesky_fn is not None):
raise ValueError(
'At most one of `marginal_fn` and `cholesky_fn` should be set.'
)
if marginal_fn is None:
if cholesky_fn is None:
self._cholesky_fn = cholesky_util.make_cholesky_with_jitter_fn(
jitter)
self._marginal_fn = make_cholesky_factored_marginal_fn(
self._cholesky_fn)
else:
self._marginal_fn = marginal_fn
self._always_yield_multivariate_normal = always_yield_multivariate_normal
with tf.name_scope('init'):
super(GaussianProcess, self).__init__(
dtype=dtype,
reparameterization_type=reparameterization.
FULLY_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name)
def __init__(self,
bin_widths,
bin_heights,
knot_slopes,
range_min=-1,
validate_args=False,
name=None):
"""Construct a new RationalQuadraticSpline bijector.
For each argument, the innermost axis indexes bins/knots and batch axes
index axes of `x`/`y` spaces. A `RationalQuadraticSpline` with a separate
transform for each of three dimensions might have `bin_widths` shaped
`[3, 32]`. To use the same spline for each of `x`'s three dimensions we may
broadcast against `x` and use a `bin_widths` parameter shaped `[32]`.
Parameters will be broadcast against each other and against the input
`x`/`y`s, so if we want fixed slopes, we can use kwarg `knot_slopes=1`.
A typical recipe for acquiring compatible bin widths and heights would be:
```python
nbins = unconstrained_vector.shape[-1]
range_min, range_max, min_bin_size = -1, 1, 1e-2
scale = range_max - range_min - nbins * min_bin_size
bin_widths = tf.math.softmax(unconstrained_vector) * scale + min_bin_size
```
Args:
bin_widths: The widths of the spans between subsequent knot `x` positions,
a floating point `Tensor`. Must be positive, and at least 1-D. Innermost
axis must sum to the same value as `bin_heights`. The knot `x` positions
will be a first at `range_min`, followed by knots at `range_min +
cumsum(bin_widths, axis=-1)`.
bin_heights: The heights of the spans between subsequent knot `y`
positions, a floating point `Tensor`. Must be positive, and at least
1-D. Innermost axis must sum to the same value as `bin_widths`. The knot
`y` positions will be a first at `range_min`, followed by knots at
`range_min + cumsum(bin_heights, axis=-1)`.
knot_slopes: The slope of the spline at each knot, a floating point
`Tensor`. Must be positive. `1`s are implicitly padded for the first and
last implicit knots corresponding to `range_min` and `range_min +
sum(bin_widths, axis=-1)`. Innermost axis size should be 1 less than
that of `bin_widths`/`bin_heights`, or 1 for broadcasting.
range_min: The `x`/`y` position of the first knot, which has implicit
slope `1`. `range_max` is implicit, and can be computed as `range_min +
sum(bin_widths, axis=-1)`. Scalar floating point `Tensor`.
validate_args: Toggles argument validation (can hurt performance).
name: Optional name scope for associated ops. (Defaults to
`'RationalQuadraticSpline'`).
"""
parameters = dict(locals())
with tf.name_scope(name or 'RationalQuadraticSpline') as name:
dtype = dtype_util.common_dtype(
[bin_widths, bin_heights, knot_slopes, range_min],
dtype_hint=tf.float32)
self._bin_widths = tensor_util.convert_nonref_to_tensor(
bin_widths, dtype=dtype, name='bin_widths')
self._bin_heights = tensor_util.convert_nonref_to_tensor(
bin_heights, dtype=dtype, name='bin_heights')
self._knot_slopes = tensor_util.convert_nonref_to_tensor(
knot_slopes, dtype=dtype, name='knot_slopes')
self._range_min = tensor_util.convert_nonref_to_tensor(
range_min, dtype=dtype, name='range_min')
super(RationalQuadraticSpline,
self).__init__(dtype=dtype,
forward_min_event_ndims=0,
validate_args=validate_args,
parameters=parameters,
name=name)
def __init__(self,
loc,
scale,
quadrature_size=8,
quadrature_fn=quadrature_scheme_lognormal_quantiles,
validate_args=False,
allow_nan_stats=True,
name='PoissonLogNormalQuadratureCompound'):
"""Constructs the PoissonLogNormalQuadratureCompound`.
Note: `probs` returned by (optional) `quadrature_fn` are presumed to be
either a length-`quadrature_size` vector or a batch of vectors in 1-to-1
correspondence with the returned `grid`. (I.e., broadcasting is only
partially supported.)
Args:
loc: `float`-like (batch of) scalar `Tensor`; the location parameter of
the LogNormal prior.
scale: `float`-like (batch of) scalar `Tensor`; the scale parameter of
the LogNormal prior.
quadrature_size: Python `int` scalar representing the number of quadrature
points.
quadrature_fn: Python callable taking `loc`, `scale`,
`quadrature_size`, `validate_args` and returning `tuple(grid, probs)`
representing the LogNormal grid and corresponding normalized weight.
Default value: `quadrature_scheme_lognormal_quantiles`.
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.
Raises:
TypeError: if `quadrature_grid` and `quadrature_probs` have different base
`dtype`.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([loc, scale], tf.float32)
self._loc = tensor_util.convert_nonref_to_tensor(loc,
name='loc',
dtype=dtype)
self._scale = tensor_util.convert_nonref_to_tensor(scale,
name='scale',
dtype=dtype)
self._quadrature_fn = quadrature_fn
dtype_util.assert_same_float_dtype([self._loc, self._scale])
self._quadrature_size = quadrature_size
super(PoissonLogNormalQuadratureCompound, self).__init__(
dtype=dtype,
reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name)
def __init__(self,
paddings=((0, 1), ),
mode='CONSTANT',
constant_values=0,
axis=None,
validate_args=False,
name=None):
"""Initializes the `Pad` bijector.
Args:
paddings: A vector-shaped `Tensor` of `int` pairs representing the number
of elements to pad on the left and right, respectively.
Default value: `((0, 1),)`.
mode: One of `'CONSTANT'`, `'REFLECT'`, or `'SYMMETRIC'`
(case-insensitive). For more details, see `tf.pad`.
constant_values: In "CONSTANT" mode, the scalar pad value to use. Must be
same type as `tensor`. For more details, see `tf.pad`.
axis: The dimensions for which `paddings` are applied. Must be 1:1 with
`paddings` or `None`.
Default value: `None` (i.e., `tf.range(start=-len(paddings), limit=0)`).
validate_args: Python `bool` indicating whether arguments should
be checked for correctness.
Default value: `False`.
name: Python `str`, name given to ops managed by this object.
Default value: `None` (i.e., `'pad'`).
"""
parameters = dict(locals())
with tf.name_scope(name or 'pad') as name:
paddings = tensor_util.convert_nonref_to_tensor(
paddings,
dtype_hint=tf.int32,
name='paddings',
as_shape_tensor=True)
if axis is None:
axis = ps.range(start=-ps.size0(paddings),
limit=0,
dtype=tf.int32,
name='axis')
else:
axis = tensor_util.convert_nonref_to_tensor(
axis,
dtype_hint=tf.int32,
name='axis',
as_shape_tensor=True)
axis_ = tf.get_static_value(axis)
if axis_ is None:
raise NotImplementedError(
'Argument `axis` must be known statically. If you need this '
'feature, please contact `[email protected]`.')
self._axis = axis
self._paddings = paddings
self._mode = mode
self._constant_values = tensor_util.convert_nonref_to_tensor(
constant_values, dtype_hint=tf.float32, name='constant_values')
min_event_ndims_ = int(-np.min(
np.pad(np.reshape(axis_, newshape=[-1]),
mode='constant',
pad_width=[[0, 1]])))
super(Pad, self).__init__(forward_min_event_ndims=min_event_ndims_,
inverse_min_event_ndims=min_event_ndims_,
is_constant_jacobian=True,
validate_args=validate_args,
parameters=parameters,
name=name)
def __init__(self,
loc=None,
scale_tril=None,
validate_args=False,
allow_nan_stats=True,
name='MultivariateNormalTriL'):
"""Construct Multivariate Normal distribution on `R^k`.
The `batch_shape` is the broadcast shape between `loc` and `scale`
arguments.
The `event_shape` is given by last dimension of the matrix implied by
`scale`. The last dimension of `loc` (if provided) must broadcast with this.
Recall that `covariance = scale @ scale.T`. A (non-batch) `scale` matrix is:
```none
scale = scale_tril
```
where `scale_tril` is lower-triangular `k x k` matrix with non-zero
diagonal, i.e., `tf.diag_part(scale_tril) != 0`.
Additional leading dimensions (if any) will index batches.
Args:
loc: Floating-point `Tensor`. If this is set to `None`, `loc` is
implicitly `0`. When specified, may have shape `[B1, ..., Bb, k]` where
`b >= 0` and `k` is the event size.
scale_tril: Floating-point, lower-triangular `Tensor` with non-zero
diagonal elements. `scale_tril` has shape `[B1, ..., Bb, k, k]` where
`b >= 0` and `k` is the event size.
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.
Raises:
ValueError: if neither `loc` nor `scale_tril` are specified.
"""
parameters = dict(locals())
if loc is None and scale_tril is None:
raise ValueError(
'Must specify one or both of `loc`, `scale_tril`.')
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([loc, scale_tril], tf.float32)
loc = tensor_util.convert_nonref_to_tensor(loc,
name='loc',
dtype=dtype)
scale_tril = tensor_util.convert_nonref_to_tensor(
scale_tril, name='scale_tril', dtype=dtype)
self._scale_tril = scale_tril
if scale_tril is None:
scale = tf.linalg.LinearOperatorIdentity(
num_rows=distribution_util.dimension_size(loc, -1),
dtype=loc.dtype,
is_self_adjoint=True,
is_positive_definite=True,
assert_proper_shapes=validate_args)
else:
# No need to validate that scale_tril is non-singular.
# LinearOperatorLowerTriangular has an assert_non_singular
# method that is called by the Bijector.
scale = tf.linalg.LinearOperatorLowerTriangular(
scale_tril,
is_non_singular=True,
is_self_adjoint=False,
is_positive_definite=False)
super(MultivariateNormalTriL,
self).__init__(loc=loc,
scale=scale,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
name=name)
self._parameters = parameters
def __init__(self,
outcomes,
logits=None,
probs=None,
rtol=None,
atol=None,
validate_args=False,
allow_nan_stats=True,
name='FiniteDiscrete'):
"""Construct a finite discrete contribution.
Args:
outcomes: A 1-D floating or integer `Tensor`, representing a list of
possible outcomes in strictly ascending order.
logits: A floating N-D `Tensor`, `N >= 1`, representing the log
probabilities of a set of FiniteDiscrete distributions. The first `N -
1` dimensions index into a batch of independent distributions and the
last dimension represents a vector of logits for each discrete value.
Only one of `logits` or `probs` should be passed in.
probs: A floating N-D `Tensor`, `N >= 1`, representing the probabilities
of a set of FiniteDiscrete distributions. The first `N - 1` dimensions
index into a batch of independent distributions and the last dimension
represents a vector of probabilities for each discrete value. Only one
of `logits` or `probs` should be passed in.
rtol: `Tensor` with same `dtype` as `outcomes`. The relative tolerance for
floating number comparison. Only effective when `outcomes` is a floating
`Tensor`. Default is `10 * eps`.
atol: `Tensor` with same `dtype` as `outcomes`. The absolute tolerance for
floating number comparison. Only effective when `outcomes` is a floating
`Tensor`. Default is `10 * eps`.
validate_args: Python `bool`, default `False`. When `True` distribution
parameters are checked for validity despite possibly degrading runtime
performance. When `False` invalid inputs may 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:
outcomes_dtype = dtype_util.common_dtype([outcomes],
dtype_hint=tf.float32)
self._outcomes = tensor_util.convert_nonref_to_tensor(
outcomes, dtype_hint=outcomes_dtype, name='outcomes')
if dtype_util.is_floating(self._outcomes.dtype):
eps = np.finfo(dtype_util.as_numpy_dtype(outcomes_dtype)).eps
self._rtol = 10 * eps if rtol is None else rtol
self._atol = 10 * eps if atol is None else atol
else:
self._rtol = None
self._atol = None
self._categorical = categorical.Categorical(
logits=logits,
probs=probs,
dtype=tf.int32,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats)
super(FiniteDiscrete, self).__init__(
dtype=self._outcomes.dtype,
reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name)
def __init__(self,
event_shape_out,
event_shape_in=(-1, ),
validate_args=False,
name=None):
"""Creates a `Reshape` bijector.
Args:
event_shape_out: An `int`-like vector-shaped `Tensor`
representing the event shape of the transformed output.
event_shape_in: An optional `int`-like vector-shape `Tensor`
representing the event shape of the input. This is required in
order to define inverse operations; the default of (-1,)
assumes a vector-shaped input.
validate_args: Python `bool` indicating whether arguments should
be checked for correctness.
name: Python `str`, name given to ops managed by this object.
Raises:
TypeError: if either `event_shape_in` or `event_shape_out` has
non-integer `dtype`.
ValueError: if either of `event_shape_in` or `event_shape_out`
has non-vector shape (`rank > 1`), or if their sizes do not
match.
"""
parameters = dict(locals())
with tf.name_scope(name or 'reshape') as name:
dtype = dtype_util.common_dtype([event_shape_out, event_shape_in],
dtype_hint=tf.int32)
event_shape_out = tensor_util.convert_nonref_to_tensor(
event_shape_out,
name='event_shape_out',
dtype=dtype,
as_shape_tensor=True)
event_shape_in = tensor_util.convert_nonref_to_tensor(
event_shape_in,
name='event_shape_in',
dtype=dtype,
as_shape_tensor=True)
forward_min_event_ndims_ = _rank_from_shape(event_shape_in)
if forward_min_event_ndims_ is None:
raise NotImplementedError(
'The length of `event_shape_in` must be statically known. For '
'dynamic support, please contact `[email protected]`.'
)
inverse_min_event_ndims_ = _rank_from_shape(event_shape_out)
if inverse_min_event_ndims_ is None:
raise NotImplementedError(
'The length of `event_shape_out` must be statically known. For '
'dynamic support, please contact `[email protected]`.'
)
self._event_shape_in = event_shape_in
self._event_shape_out = event_shape_out
super(Reshape, self).__init__(
forward_min_event_ndims=forward_min_event_ndims_,
inverse_min_event_ndims=inverse_min_event_ndims_,
is_constant_jacobian=True,
validate_args=validate_args,
parameters=parameters,
name=name or 'reshape')
def __init__(self,
perm=None,
rightmost_transposed_ndims=None,
validate_args=False,
name='transpose'):
"""Instantiates the `Transpose` bijector.
Args:
perm: Positive `int32` vector-shaped `Tensor` representing permutation of
rightmost dims (for forward transformation). Note that the `0`th index
represents the first of the rightmost dims and the largest value must be
`rightmost_transposed_ndims - 1` and corresponds to `tf.rank(x) - 1`.
Only one of `perm` and `rightmost_transposed_ndims` can (and must) be
specified. The number of elements in a permutation must have a value
that can be determined statically.
Default value:
`tf.range(start=rightmost_transposed_ndims, limit=-1, delta=-1)`.
rightmost_transposed_ndims: Positive `int32` scalar-shaped `Tensor`
representing the number of rightmost dimensions to permute.
Only one of `perm` and `rightmost_transposed_ndims` can (and must) be
specified. If `rightmost_transposed_ndims` is specified, the rightmost
dims are reversed. This argument must have a value that can be
determined statically.
Default value: `tf.size(perm)`.
validate_args: Python `bool` indicating whether arguments should be
checked for correctness.
name: Python `str` name given to ops managed by this object.
Raises:
ValueError: if both or neither `perm` and `rightmost_transposed_ndims` are
specified.
NotImplementedError: if `rightmost_transposed_ndims` is not known prior to
graph execution.
"""
with tf.name_scope(name) as name:
# We need to determine `forward_min_event_ndims` statically, which
# requires that we know `rightmost_transposed_ndims` statically.
# So the corresponding assertions go here rather than in
# `_parameter_control_dependencies`
if (rightmost_transposed_ndims is None) == (perm is None):
raise ValueError('Must specify exactly one of '
'`rightmost_transposed_ndims` and `perm`.')
if rightmost_transposed_ndims is not None:
rightmost_transposed_ndims = tensor_util.convert_nonref_to_tensor(
rightmost_transposed_ndims, dtype_hint=np.int32)
if not dtype_util.is_integer(rightmost_transposed_ndims.dtype):
raise TypeError(
'`rightmost_transposed_ndims` must be integer type.')
rightmost_transposed_ndims_ = tf.get_static_value(
rightmost_transposed_ndims)
if rightmost_transposed_ndims_ is None:
raise NotImplementedError(
'`rightmost_transposed_ndims` must be '
'known prior to graph execution.')
msg = '`rightmost_transposed_ndims` must be non-negative.'
if rightmost_transposed_ndims_ < 0:
raise ValueError(
msg[:-1] +
', saw: {}.'.format(rightmost_transposed_ndims_))
perm_start = (distribution_util.prefer_static_value(
rightmost_transposed_ndims) - 1)
perm = tf.range(start=perm_start,
limit=-1,
delta=-1,
name='perm')
else: # perm is not None:
perm = tensor_util.convert_nonref_to_tensor(
perm, dtype_hint=np.int32, name='perm')
rightmost_transposed_ndims = tf.size(
perm, name='rightmost_transposed_ndims')
rightmost_transposed_ndims_ = tf.get_static_value(
rightmost_transposed_ndims)
# TODO(b/110828604): If bijector base class ever supports dynamic
# `min_event_ndims`, then this class already works dynamically and the
# following five lines can be removed.
if rightmost_transposed_ndims_ is None:
raise NotImplementedError(
'`rightmost_transposed_ndims` must be '
'known prior to graph execution.')
else:
rightmost_transposed_ndims_ = int(rightmost_transposed_ndims_)
self._perm = perm
self._rightmost_transposed_ndims = rightmost_transposed_ndims
self._initial_rightmost_transposed_ndims = rightmost_transposed_ndims_
super(Transpose, self).__init__(
forward_min_event_ndims=rightmost_transposed_ndims_,
is_constant_jacobian=True,
validate_args=validate_args,
name=name)
def __init__(self,
kernel,
x1,
x2=None,
num_matmul_parts=None,
is_non_singular=None,
is_self_adjoint=None,
is_positive_definite=None,
is_square=None,
name=None):
"""Initializes the LinearOperator.
This object implicitly represents the covariance matrix of `x1` and `x2`
(`x1` if `x2` not provided) under the given `kernel`. This operator assumes
one example dimension on each set of index points, which indexes the
corresponding axis of the matrix. All outer dimensions are considered
batch dimensions.
Use this to avoid materializing the full matrix for such operations as:
- accessing diagonal (`diag_part` method)
- accessing a [batch of] row[s] (`row` method)
- accessing a [batch of] column[s] (`col` method)
Use this to perform matrix-vector products on very large covariance matrices
by chunking the covariance matrix into parts, computing just that part of
the output, computing a part of the covariance matrix, and the matrix-vector
product, then forgetting the partial covariance matrix. Internally, uses
recomputed gradients to avoid retaining infeasibly-large intermediate
activations.
Args:
kernel: A `tfp.math.psd_kernels.PositiveSemidefiniteKernel` instance.
x1: A floating point `Tensor`, the row index points.
x2: Optional, a floating point `Tensor`, the column index points. If not
provided, uses `x1`.
num_matmul_parts: An optional Python `int` specifying the number of
partitions into which the matrix should be broken when applying this
linear operator. (An extra remainder partition is implied for uneven
partitioning.) Because the use-case is avoiding a memory blowup, the
partitioned matmul uses `parallel_iterations=1` and `back_prop=False`.
is_non_singular: Expect that this operator is non-singular.
is_self_adjoint: Expect that this operator is equal to its hermitian
transpose. If `dtype` is real, this is equivalent to being symmetric.
is_positive_definite: Expect that this operator is positive definite,
meaning the quadratic form `x^H A x` has positive real part for all
nonzero `x`. Note that we do not require the operator to be
self-adjoint to be positive-definite. See:
https://en.wikipedia.org/wiki/Positive-definite_matrix#Extension_for_non-symmetric_matrices
is_square: Expect that this operator acts like square [batch] matrices.
name: Optional name for related ops.
"""
dtype = dtype_util.common_dtype([kernel, x1, x2],
dtype_hint=tf.float64)
self._kernel = kernel
self._x1 = tensor_util.convert_nonref_to_tensor(x1, dtype=dtype)
self._x2 = tensor_util.convert_nonref_to_tensor(x2, dtype=dtype)
self._num_matmul_parts = tensor_util.convert_nonref_to_tensor(
num_matmul_parts, dtype=tf.int32)
if self._x2 is None:
if is_non_singular is not None and not is_non_singular:
raise ValueError(
'Operator is non-singular with a single set of index points.'
)
is_non_singular = True
if is_self_adjoint is not None and not is_self_adjoint:
raise ValueError(
'Operator is self-adjoint with a single set of index points.'
)
is_self_adjoint = True
if is_positive_definite is not None and not is_positive_definite:
raise ValueError(
'Operator is positive-definite with a single set of index points.'
)
is_positive_definite = True
if is_square is not None and not is_square:
raise ValueError(
'Operator is square with a single set of index points.')
is_square = True
super(LinearOperatorPSDKernel,
self).__init__(dtype,
is_non_singular=is_non_singular,
is_self_adjoint=is_self_adjoint,
is_positive_definite=is_positive_definite,
is_square=is_square,
name=name or 'LinearOperatorPSDKernel')
def __init__(self, x, name=None):
super(FakeModule, self).__init__(name)
self.x = tensor_util.convert_nonref_to_tensor(x)
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 = 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 Gumbel bijector.
self._gumbel_bijector = gumbel_cdf_bijector.GumbelCDF(
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=dtype),
allow_nan_stats=allow_nan_stats),
# The Gumbel bijector encodes the CDF function as the forward,
# and hence needs to be inverted.
bijector=invert_bijector.Invert(self._gumbel_bijector,
validate_args=validate_args),
parameters=parameters,
name=name)
def test_tf_tensor(self):
x = tf.constant(1.)
y = tensor_util.convert_nonref_to_tensor(x)
self.assertIs(x, y)
def __init__(self,
df,
scale=None,
scale_tril=None,
input_output_cholesky=False,
validate_args=False,
allow_nan_stats=True,
name='Wishart'):
"""Construct Wishart distributions.
Args:
df: `float` or `double` `Tensor`. Degrees of freedom, must be greater than
or equal to dimension of the scale matrix.
scale: `float` or `double` `Tensor`. The symmetric positive definite
scale matrix of the distribution. Exactly one of `scale` and
'scale_tril` must be passed.
scale_tril: `float` or `double` `Tensor`. The Cholesky factorization
of the symmetric positive definite scale matrix of the distribution.
Exactly one of `scale` and 'scale_tril` must be passed.
input_output_cholesky: Python `bool`. If `True`, functions whose input or
output have the semantics of samples assume inputs are in Cholesky form
and return outputs in Cholesky form. In particular, if this flag is
`True`, input to `log_prob` is presumed of Cholesky form and output from
`sample`, `mean`, and `mode` are of Cholesky form. Setting this
argument to `True` is purely a computational optimization and does not
change the underlying distribution; for instance, `mean` returns the
Cholesky of the mean, not the mean of Cholesky factors. The `variance`
and `stddev` methods are unaffected by this flag.
Default value: `False` (i.e., input/output does not have Cholesky
semantics).
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.
Raises:
ValueError: if zero or both of 'scale' and 'scale_tril' are passed in.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
if (scale is None) == (scale_tril is None):
raise ValueError(
'Must pass scale or scale_tril, but not both.')
dtype = dtype_util.common_dtype([df, scale, scale_tril],
tf.float32)
df = tensor_util.convert_nonref_to_tensor(df,
name='df',
dtype=dtype)
if scale is not None:
scale = tensor_util.convert_nonref_to_tensor(scale,
name='scale',
dtype=dtype)
self._scale_tril = tf.linalg.cholesky(scale)
else: # scale_tril is not None
self._scale_tril = tensor_util.convert_nonref_to_tensor(
scale_tril, name='scale_tril', dtype=dtype)
self._scale_full = scale
super(Wishart,
self).__init__(df=df,
scale=tf.linalg.LinearOperatorLowerTriangular(
tril=self._scale_tril,
is_non_singular=True,
is_positive_definite=True,
is_square=True),
input_output_cholesky=input_output_cholesky,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
name=name)
self._parameters = parameters
def __init__(self,
power,
dtype=tf.int32,
interpolate_nondiscrete=True,
sample_maximum_iterations=100,
validate_args=False,
allow_nan_stats=False,
name='Zipf'):
"""Initialize a batch of Zipf distributions.
Args:
power: `Float` like `Tensor` representing the power parameter. Must be
strictly greater than `1`.
dtype: The `dtype` of `Tensor` returned by `sample`.
Default value: `tf.int32`.
interpolate_nondiscrete: Python `bool`. When `False`, `log_prob` returns
`-inf` (and `prob` returns `0`) for non-integer inputs. When `True`,
`log_prob` evaluates the continuous function `-power log(k) -
log(zeta(power))` , which matches the Zipf pmf at integer arguments `k`
(note that this function is not itself a normalized probability
log-density).
Default value: `True`.
sample_maximum_iterations: Maximum number of iterations of allowable
iterations in `sample`. When `validate_args=True`, samples which fail to
reach convergence (subject to this cap) are masked out with
`self.dtype.min` or `nan` depending on `self.dtype.is_integer`.
Default value: `100`.
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: `False`.
name: Python `str` name prefixed to Ops created by this class.
Default value: `'Zipf'`.
Raises:
TypeError: if `power` is not `float` like.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
self._power = tensor_util.convert_nonref_to_tensor(
power,
name='power',
dtype=dtype_util.common_dtype([power], dtype_hint=tf.float32))
if (not dtype_util.is_floating(self._power.dtype) or
dtype_util.base_equal(self._power.dtype, tf.float16)):
raise TypeError(
'power.dtype ({}) is not a supported `float` type.'.format(
dtype_util.name(self._power.dtype)))
self._interpolate_nondiscrete = interpolate_nondiscrete
self._sample_maximum_iterations = sample_maximum_iterations
super(Zipf, self).__init__(
dtype=dtype,
reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name)
def __init__(self,
num_or_size_splits,
axis=-1,
validate_args=False,
name='split'):
"""Creates the bijector.
Args:
num_or_size_splits: Either a Python integer indicating the number of
splits along `axis` or a 1-D integer `Tensor` or Python list containing
the sizes of each output tensor along `axis`. If a list/`Tensor`, it may
contain at most one value of `-1`, which indicates a split size that is
unknown and determined from input.
axis: A negative integer or scalar `int32` `Tensor`. The dimension along
which to split. Must be negative to enable the bijector to support
arbitrary batch dimensions. Defaults to -1 (note that this is different
from the `tf.Split` default of `0`). Must be statically known.
validate_args: Python `bool` indicating whether arguments should
be checked for correctness.
name: Python `str`, name given to ops managed by this object.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
if isinstance(num_or_size_splits, numbers.Integral):
self._num_splits = num_or_size_splits
self._split_sizes = None
else:
self._split_sizes = tensor_util.convert_nonref_to_tensor(
num_or_size_splits,
name='num_or_size_splits',
dtype=tf.int32)
if tensorshape_util.rank(self._split_sizes.shape) != 1:
raise ValueError(
'`num_or_size_splits` must be an integer or 1-D `Tensor`.'
)
num_splits = tensorshape_util.as_list(
self._split_sizes.shape)[0]
if num_splits is None:
raise ValueError(
'If `num_or_size_splits` is a vector of split sizes '
'it must have a statically-known number of '
'elements.')
self._num_splits = num_splits
static_axis = tf.get_static_value(axis)
if static_axis is None:
raise ValueError('`axis` must be statically known.')
if static_axis >= 0:
raise ValueError(
'`axis` must be negative. Got {}'.format(axis))
self._axis = tf.convert_to_tensor(axis, tf.int32)
super(Split, self).__init__(forward_min_event_ndims=-axis,
inverse_min_event_ndims=[-axis] *
self.num_splits,
is_constant_jacobian=True,
validate_args=validate_args,
parameters=parameters,
name=name)
def __init__(self, x, y, name=None):
with tf.name_scope(name or 'Adder') as name:
self._x = tensor_util.convert_nonref_to_tensor(x)
self._y = tensor_util.convert_nonref_to_tensor(y)
self._name = name
def __init__(self,
rate=None,
log_rate=None,
interpolate_nondiscrete=True,
validate_args=False,
allow_nan_stats=True,
name='Poisson'):
"""Initialize a batch of Poisson distributions.
Args:
rate: Floating point tensor, the rate parameter. `rate` must be positive.
Must specify exactly one of `rate` and `log_rate`.
log_rate: Floating point tensor, the log of the rate parameter.
Must specify exactly one of `rate` and `log_rate`.
interpolate_nondiscrete: Python `bool`. When `False`,
`log_prob` returns `-inf` (and `prob` returns `0`) for non-integer
inputs. When `True`, `log_prob` evaluates the continuous function
`k * log_rate - lgamma(k+1) - rate`, which matches the Poisson pmf
at integer arguments `k` (note that this function is not itself
a normalized probability log-density).
Default value: `True`.
validate_args: Python `bool`. 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`. 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.
Raises:
ValueError: if none or both of `rate`, `log_rate` are specified.
TypeError: if `rate` is not a float-type.
TypeError: if `log_rate` is not a float-type.
"""
parameters = dict(locals())
if (rate is None) == (log_rate is None):
raise ValueError(
'Must specify exactly one of `rate` and `log_rate`.')
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([rate, log_rate],
dtype_hint=tf.float32)
if not dtype_util.is_floating(dtype):
raise TypeError(
'[log_]rate.dtype ({}) is a not a float-type.'.format(
dtype_util.name(dtype)))
self._rate = tensor_util.convert_nonref_to_tensor(rate,
name='rate',
dtype=dtype)
self._log_rate = tensor_util.convert_nonref_to_tensor(
log_rate, name='log_rate', dtype=dtype)
self._interpolate_nondiscrete = interpolate_nondiscrete
super(Poisson, self).__init__(
dtype=dtype,
reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name)
def __init__(self,
df,
kernel,
index_points=None,
mean_fn=None,
observation_noise_variance=0.,
marginal_fn=None,
jitter=1e-6,
validate_args=False,
allow_nan_stats=False,
name='StudentTProcess'):
"""Instantiate a StudentTProcess Distribution.
Args:
df: Positive Floating-point `Tensor` representing the degrees of freedom.
Must be greater than 2.
kernel: `PositiveSemidefiniteKernel`-like instance representing the
TP's covariance function.
index_points: `float` `Tensor` representing finite (batch of) vector(s) of
points in the index set over which the TP is defined. Shape has the form
`[b1, ..., bB, e, f1, ..., fF]` where `F` is the number of feature
dimensions and must equal `kernel.feature_ndims` and `e` is the number
(size) of index points in each batch. Ultimately this distribution
corresponds to a `e`-dimensional multivariate Student's T. The batch
shape must be broadcastable with `kernel.batch_shape` and any batch dims
yielded by `mean_fn`.
mean_fn: Python `callable` that acts on `index_points` to produce a (batch
of) vector(s) of mean values at `index_points`. Takes a `Tensor` of
shape `[b1, ..., bB, f1, ..., fF]` and returns a `Tensor` whose shape is
broadcastable with `[b1, ..., bB]`. Default value: `None` implies
constant zero function.
observation_noise_variance: `float` `Tensor` representing (batch of)
scalar variance(s) of the noise in the Normal likelihood
distribution of the model. If batched, the batch shape must be
broadcastable with the shapes of all other batched parameters
(`kernel.batch_shape`, `index_points`, etc.).
Default value: `0.`
marginal_fn: A Python callable that takes a location, covariance matrix,
optional `validate_args`, `allow_nan_stats` and `name` arguments, and
returns a multivariate normal subclass of `tfd.Distribution`.
Default value: `None`, in which case a Cholesky-factorizing function is
is created using `make_cholesky_factorizing_marginal_fn` and the
`jitter` argument.
jitter: `float` scalar `Tensor` added to the diagonal of the covariance
matrix to ensure positive definiteness of the covariance matrix.
Default value: `1e-6`.
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: `False`.
name: Python `str` name prefixed to Ops created by this class.
Default value: "StudentTProcess".
Raises:
ValueError: if `mean_fn` is not `None` and is not callable.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype(
[df, index_points, observation_noise_variance, jitter],
tf.float32)
df = tensor_util.convert_nonref_to_tensor(df,
dtype=dtype,
name='df')
observation_noise_variance = tensor_util.convert_nonref_to_tensor(
observation_noise_variance,
dtype=dtype,
name='observation_noise_variance')
index_points = tensor_util.convert_nonref_to_tensor(
index_points, dtype=dtype, name='index_points')
jitter = tensor_util.convert_nonref_to_tensor(jitter,
dtype=dtype,
name='jitter')
self._kernel = kernel
self._index_points = index_points
# Default to a constant zero function, borrowing the dtype from
# index_points to ensure consistency.
if mean_fn is None:
mean_fn = lambda x: tf.zeros([1], dtype=dtype)
else:
if not callable(mean_fn):
raise ValueError('`mean_fn` must be a Python callable')
self._df = df
self._observation_noise_variance = observation_noise_variance
self._mean_fn = mean_fn
self._jitter = jitter
if marginal_fn is None:
self._marginal_fn = make_cholesky_factored_marginal_fn(jitter)
else:
self._marginal_fn = marginal_fn
with tf.name_scope('init'):
super(StudentTProcess, self).__init__(
dtype=dtype,
reparameterization_type=reparameterization.
FULLY_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name)
def __init__(self,
loc,
scale,
skewness=None,
tailweight=None,
distribution=None,
validate_args=False,
allow_nan_stats=True,
name="SinhArcsinh"):
"""Construct SinhArcsinh distribution on `(-inf, inf)`.
Arguments `(loc, scale, skewness, tailweight)` must have broadcastable shape
(indexing batch dimensions). They must all have the same `dtype`.
Args:
loc: Floating-point `Tensor`.
scale: `Tensor` of same `dtype` as `loc`.
skewness: Skewness parameter. Default is `0.0` (no skew).
tailweight: Tailweight parameter. Default is `1.0` (unchanged tailweight)
distribution: `tf.Distribution`-like instance. Distribution that is
transformed to produce this distribution.
Default is `tfd.Normal(0., 1.)`.
Must be a scalar-batch, scalar-event distribution. Typically
`distribution.reparameterization_type = FULLY_REPARAMETERIZED` or it is
a function of non-trainable parameters. WARNING: If you backprop through
a `SinhArcsinh` sample and `distribution` is not
`FULLY_REPARAMETERIZED` yet is a function of trainable variables, then
the gradient will be incorrect!
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([loc, scale, skewness, tailweight],
tf.float32)
self._loc = tensor_util.convert_nonref_to_tensor(loc,
name="loc",
dtype=dtype)
self._scale = tensor_util.convert_nonref_to_tensor(scale,
name="scale",
dtype=dtype)
tailweight = 1. if tailweight is None else tailweight
has_default_skewness = skewness is None
skewness = 0. if has_default_skewness else skewness
self._tailweight = tensor_util.convert_nonref_to_tensor(
tailweight, name="tailweight", dtype=dtype)
self._skewness = tensor_util.convert_nonref_to_tensor(
skewness, name="skewness", dtype=dtype)
batch_shape = distribution_util.get_broadcast_shape(
self._loc, self._scale, self._tailweight, self._skewness)
# Recall, with Z a random variable,
# Y := loc + scale * F(Z),
# F(Z) := Sinh( (Arcsinh(Z) + skewness) * tailweight ) * C
# C := 2 / F_0(2)
# F_0(Z) := Sinh( Arcsinh(Z) * tailweight )
if distribution is None:
distribution = normal.Normal(loc=tf.zeros([], dtype=dtype),
scale=tf.ones([], dtype=dtype),
allow_nan_stats=allow_nan_stats,
validate_args=validate_args)
else:
asserts = distribution_util.maybe_check_scalar_distribution(
distribution, dtype, validate_args)
if asserts:
self._loc = distribution_util.with_dependencies(
asserts, self._loc)
# Make the SAS bijector, 'F'.
f = sinh_arcsinh_bijector.SinhArcsinh(skewness=self._skewness,
tailweight=self._tailweight,
validate_args=validate_args)
# Make the AffineScalar bijector, Z --> loc + scale * Z (2 / F_0(2))
affine = affine_scalar_bijector.AffineScalar(
shift=self._loc,
scale=self._scale,
validate_args=validate_args)
bijector = chain_bijector.Chain([affine, f])
super(SinhArcsinh, self).__init__(distribution=distribution,
bijector=bijector,
batch_shape=batch_shape,
validate_args=validate_args,
name=name)
self._parameters = parameters
