def __init__(self,
loc,
concentration,
validate_args=False,
allow_nan_stats=True,
name='InverseGaussian'):
"""Constructs inverse Gaussian distribution with `loc` and `concentration`.
Args:
loc: Floating-point `Tensor`, the loc params. Must contain only positive
values.
concentration: Floating-point `Tensor`, the concentration params.
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` (i.e. do not validate args).
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: 'InverseGaussian'.
"""
parameters = dict(locals())
with tf.name_scope(name):
dtype = dtype_util.common_dtype([loc, concentration],
dtype_hint=tf.float32)
self._concentration = tensor_util.convert_nonref_to_tensor(
concentration, dtype=dtype, name='concentration')
self._loc = tensor_util.convert_nonref_to_tensor(loc,
dtype=dtype,
name='loc')
super(InverseGaussian, self).__init__(
dtype=self._loc.dtype,
reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name)
def __init__(self,
low=0.,
high=1.,
validate_args=False,
allow_nan_stats=True,
name='Uniform'):
"""Initialize a batch of Uniform distributions.
Args:
low: Floating point tensor, lower boundary of the output interval. Must
have `low < high`.
high: Floating point tensor, upper boundary of the output interval. Must
have `low < high`.
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:
InvalidArgumentError: if `low >= high` and `validate_args=False`.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([low, high], tf.float32)
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)
super(Uniform,
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,
concentration1,
concentration0,
validate_args=False,
allow_nan_stats=True,
name="Beta"):
"""Initialize a batch of Beta distributions.
Args:
concentration1: Positive floating-point `Tensor` indicating mean
number of successes; aka "alpha". Implies `self.dtype` and
`self.batch_shape`, i.e.,
`concentration1.shape = [N1, N2, ..., Nm] = self.batch_shape`.
concentration0: Positive floating-point `Tensor` indicating mean
number of failures; aka "beta". Otherwise has same semantics as
`concentration1`.
validate_args: Python `bool`, default `False`. When `True` distribution
parameters are checked for validity despite possibly degrading runtime
performance. When `False` invalid inputs may silently render incorrect
outputs.
allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
(e.g., mean, mode, variance) use the value "`NaN`" to indicate the
result is undefined. When `False`, an exception is raised if one or
more of the statistic's batch members are undefined.
name: Python `str` name prefixed to Ops created by this class.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([concentration1, concentration0],
dtype_hint=tf.float32)
self._concentration1 = tensor_util.convert_nonref_to_tensor(
concentration1, dtype=dtype, name="concentration1")
self._concentration0 = tensor_util.convert_nonref_to_tensor(
concentration0, dtype=dtype, name="concentration0")
super(Beta,
self).__init__(dtype=dtype,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
reparameterization_type=reparameterization.
FULLY_REPARAMETERIZED,
parameters=parameters,
name=name)
def __init__(self,
loc,
scale,
validate_args=False,
allow_nan_stats=True,
name='Logistic'):
"""Construct Logistic distributions with mean 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). 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.
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: The name to give Ops created by the initializer.
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)
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)
super(Logistic, self).__init__(
dtype=self._scale.dtype,
reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name)
def __init__(self,
concentration,
scale=1.,
validate_args=False,
allow_nan_stats=True,
name='Pareto'):
"""Construct Pareto distribution with `concentration` and `scale`.
Args:
concentration: Floating point tensor. Must contain only positive values.
scale: Floating point tensor, equivalent to `mode`. `scale` also
restricts the domain of this distribution to be in `[scale, inf)`.
Must contain only positive values. Default value: `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` (i.e. do not validate args).
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: 'Pareto'.
"""
parameters = dict(locals())
with tf.name_scope(name):
dtype = dtype_util.common_dtype([concentration, scale],
dtype_hint=tf.float32)
self._concentration = tensor_util.convert_nonref_to_tensor(
concentration, name='concentration', dtype=dtype)
self._scale = tensor_util.convert_nonref_to_tensor(scale,
name='scale',
dtype=dtype)
super(Pareto,
self).__init__(dtype=self._concentration.dtype,
reparameterization_type=reparameterization.
FULLY_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name)
def _setdiff1d(a, b, aminusb=True, validate_indices=True):
"""Compute set difference of elements in last dimension of `a` and `b`."""
if not aminusb:
raise NotImplementedError(
'Argument `aminusb != True` is currently unimplemented.')
if not validate_indices:
raise NotImplementedError(
'Argument `validate_indices != True` is currently unimplemented.')
with tf.name_scope('setdiff1d'):
dtype = dtype_util.as_numpy_dtype(
dtype_util.common_dtype([a, b], dtype_hint=tf.int32))
a_ = tf.get_static_value(a)
b_ = tf.get_static_value(b)
if a_ is None or b_ is None:
a = tf.convert_to_tensor(a, dtype=dtype, name='a')
b = tf.convert_to_tensor(b, dtype=dtype, name='b')
return tf.sparse.to_dense(tf.sets.difference(
a[tf.newaxis], b[tf.newaxis]))[0]
a_ = np.array(a_, dtype=dtype)
b_ = np.array(b_, dtype=dtype)
return np.setdiff1d(a_, b_)
def __init__(self,
scale,
validate_args=False,
allow_nan_stats=True,
name='Horseshoe'):
"""Construct a Horseshoe distribution with `scale`.
Args:
scale: Floating point tensor; the scales of the distribution(s).
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` (i.e., do not validate args).
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: 'Horseshoe'.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([scale], dtype_hint=tf.float32)
self._scale = tensor_util.convert_nonref_to_tensor(scale,
name='scale',
dtype=dtype)
self._half_cauchy = half_cauchy.HalfCauchy(
loc=tf.zeros([], dtype=dtype),
scale=tf.ones([], dtype=dtype),
allow_nan_stats=True)
super(Horseshoe,
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,
map_values,
validate_args=False,
name='categorical_to_discrete'):
"""Instantiates `CategoricalToDiscrete` bijector.
Args:
map_values: 1D numerical tensor of discrete values to map to, sorted in
strictly increasing order.
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([map_values], tf.float32)
self._map_values = tensor_util.convert_nonref_to_tensor(
map_values, name='map_values', dtype=dtype)
super(CategoricalToDiscrete,
self).__init__(forward_min_event_ndims=0,
is_constant_jacobian=True,
validate_args=validate_args,
name=name)
def __init__(self,
rate,
validate_args=False,
allow_nan_stats=True,
name="Exponential"):
"""Construct Exponential distribution with parameter `rate`.
Args:
rate: Floating point tensor, equivalent to `1 / mean`. 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.
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())
# Even though all statistics of are defined for valid inputs, this is not
# true in the parent class "Gamma." Therefore, passing
# allow_nan_stats=True
# through to the parent class results in unnecessary asserts.
with tf.name_scope(name) as name:
self._rate = tensor_util.convert_nonref_to_tensor(
rate,
name="rate",
dtype=dtype_util.common_dtype([rate], dtype_hint=tf.float32))
super(Exponential, self).__init__(concentration=1.,
rate=self._rate,
allow_nan_stats=allow_nan_stats,
validate_args=validate_args,
name=name)
self._parameters = parameters
def __init__(self,
df,
validate_args=False,
allow_nan_stats=True,
name='Chi'):
"""Construct Chi distributions with parameter `df`.
Args:
df: Floating point tensor, the degrees of freedom of the
distribution(s). `df` 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.
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.
Default value: `'Chi'`.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([df], dtype_hint=tf.float32)
self._df = tensor_util.convert_nonref_to_tensor(df,
name='df',
dtype=dtype)
super(Chi, self).__init__(
distribution=chi2.Chi2(df=self._df,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats),
bijector=invert_bijector.Invert(
square_bijector.Square(validate_args=validate_args)),
validate_args=validate_args,
parameters=parameters,
name=name)
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:
with tf.name_scope("init"):
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 = tf.convert_to_tensor(df, name="df", dtype=dtype)
if scale is not None:
scale = tf.convert_to_tensor(scale,
name="scale",
dtype=dtype)
if validate_args:
scale = distribution_util.assert_symmetric(scale)
scale_tril = tf.linalg.cholesky(scale)
else: # scale_tril is not None
scale_tril = tf.convert_to_tensor(scale_tril,
name="scale_tril",
dtype=dtype)
if validate_args:
scale_tril = distribution_util.with_dependencies([
assert_util.assert_positive(
tf.linalg.diag_part(scale_tril),
message="scale_tril must be positive definite"
),
assert_util.assert_equal(
tf.shape(scale_tril)[-1],
tf.shape(scale_tril)[-2],
message="scale_tril must be square")
], scale_tril)
super(Wishart, self).__init__(
df=df,
scale_operator=tf.linalg.LinearOperatorLowerTriangular(
tril=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,
loc=None,
scale_diag=None,
scale_identity_multiplier=None,
scale_perturb_factor=None,
scale_perturb_diag=None,
validate_args=False,
allow_nan_stats=True,
name="MultivariateNormalDiagPlusLowRank"):
"""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 = diag(scale_diag + scale_identity_multiplier ones(k)) +
scale_perturb_factor @ diag(scale_perturb_diag) @ scale_perturb_factor.T
```
where:
* `scale_diag.shape = [k]`,
* `scale_identity_multiplier.shape = []`,
* `scale_perturb_factor.shape = [k, r]`, typically `k >> r`, and,
* `scale_perturb_diag.shape = [r]`.
Additional leading dimensions (if any) will index batches.
If both `scale_diag` and `scale_identity_multiplier` are `None`, then
`scale` is the Identity matrix.
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_diag: Non-zero, floating-point `Tensor` representing a diagonal
matrix added to `scale`. May have shape `[B1, ..., Bb, k]`, `b >= 0`,
and characterizes `b`-batches of `k x k` diagonal matrices added to
`scale`. When both `scale_identity_multiplier` and `scale_diag` are
`None` then `scale` is the `Identity`.
scale_identity_multiplier: Non-zero, floating-point `Tensor` representing
a scaled-identity-matrix added to `scale`. May have shape
`[B1, ..., Bb]`, `b >= 0`, and characterizes `b`-batches of scaled
`k x k` identity matrices added to `scale`. When both
`scale_identity_multiplier` and `scale_diag` are `None` then `scale` is
the `Identity`.
scale_perturb_factor: Floating-point `Tensor` representing a rank-`r`
perturbation added to `scale`. May have shape `[B1, ..., Bb, k, r]`,
`b >= 0`, and characterizes `b`-batches of rank-`r` updates to `scale`.
When `None`, no rank-`r` update is added to `scale`.
scale_perturb_diag: Floating-point `Tensor` representing a diagonal matrix
inside the rank-`r` perturbation added to `scale`. May have shape
`[B1, ..., Bb, r]`, `b >= 0`, and characterizes `b`-batches of `r x r`
diagonal matrices inside the perturbation added to `scale`. When
`None`, an identity matrix is used inside the perturbation. Can only be
specified if `scale_perturb_factor` is also specified.
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 at most `scale_identity_multiplier` is specified.
"""
parameters = dict(locals())
def _convert_to_tensor(x, name, dtype=None):
return None if x is None else tf.convert_to_tensor(
x, name=name, dtype=dtype)
with tf.name_scope(name) as name:
with tf.name_scope("init"):
dtype = dtype_util.common_dtype([
loc, scale_diag, scale_identity_multiplier, scale_perturb_factor,
scale_perturb_diag
], tf.float32)
has_low_rank = (scale_perturb_factor is not None or
scale_perturb_diag is not None)
scale = distribution_util.make_diag_scale(
loc=loc,
scale_diag=scale_diag,
scale_identity_multiplier=scale_identity_multiplier,
validate_args=validate_args,
assert_positive=has_low_rank,
dtype=dtype)
scale_perturb_factor = _convert_to_tensor(
scale_perturb_factor, name="scale_perturb_factor", dtype=dtype)
scale_perturb_diag = _convert_to_tensor(
scale_perturb_diag, name="scale_perturb_diag", dtype=dtype)
if has_low_rank:
scale = tf.linalg.LinearOperatorLowRankUpdate(
scale,
u=scale_perturb_factor,
diag_update=scale_perturb_diag,
is_diag_update_positive=scale_perturb_diag is None,
is_non_singular=True, # Implied by is_positive_definite=True.
is_self_adjoint=True,
is_positive_definite=True,
is_square=True)
super(MultivariateNormalDiagPlusLowRank, 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,
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,
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_bijector.Gumbel(
loc=loc, scale=scale, validate_args=validate_args)
# Because the uniform sampler generates samples in `[0, 1)` this would
# cause samples to lie in `(inf, -inf]` instead of `(inf, -inf)`. To fix
# this, we use `np.finfo(dtype_util.as_numpy_dtype(self.dtype).tiny`
# because it is the smallest, positive, 'normal' number.
super(Gumbel, self).__init__(
distribution=uniform.Uniform(low=np.finfo(
dtype_util.as_numpy_dtype(dtype)).tiny,
high=tf.ones([], dtype=dtype),
allow_nan_stats=allow_nan_stats),
# The Gumbel bijector encodes the quantile function as the forward,
# and hence needs to be inverted.
bijector=invert_bijector.Invert(self._gumbel_bijector),
batch_shape=distribution_util.get_broadcast_shape(loc, scale),
parameters=parameters,
name=name)
def __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)
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 pivoted_cholesky(matrix, max_rank, diag_rtol=1e-3, name=None):
"""Computes the (partial) pivoted cholesky decomposition of `matrix`.
The pivoted Cholesky is a low rank approximation of the Cholesky decomposition
of `matrix`, i.e. as described in [(Harbrecht et al., 2012)][1]. The
currently-worst-approximated diagonal element is selected as the pivot at each
iteration. This yields from a `[B1...Bn, N, N]` shaped `matrix` a `[B1...Bn,
N, K]` shaped rank-`K` approximation `lr` such that `lr @ lr.T ~= matrix`.
Note that, unlike the Cholesky decomposition, `lr` is not triangular even in
a rectangular-matrix sense. However, under a permutation it could be made
triangular (it has one more zero in each column as you move to the right).
Such a matrix can be useful as a preconditioner for conjugate gradient
optimization, i.e. as in [(Wang et al. 2019)][2], as matmuls and solves can be
cheaply done via the Woodbury matrix identity, as implemented by
`tf.linalg.LinearOperatorLowRankUpdate`.
Args:
matrix: Floating point `Tensor` batch of symmetric, positive definite
matrices.
max_rank: Scalar `int` `Tensor`, the rank at which to truncate the
approximation.
diag_rtol: Scalar floating point `Tensor` (same dtype as `matrix`). If the
errors of all diagonal elements of `lr @ lr.T` are each lower than
`element * diag_rtol`, iteration is permitted to terminate early.
name: Optional name for the op.
Returns:
lr: Low rank pivoted Cholesky approximation of `matrix`.
#### References
[1]: H Harbrecht, M Peters, R Schneider. On the low-rank approximation by the
pivoted Cholesky decomposition. _Applied numerical mathematics_,
62(4):428-440, 2012.
[2]: K. A. Wang et al. Exact Gaussian Processes on a Million Data Points.
_arXiv preprint arXiv:1903.08114_, 2019. https://arxiv.org/abs/1903.08114
"""
with tf.name_scope(name or 'pivoted_cholesky'):
dtype = dtype_util.common_dtype([matrix, diag_rtol],
dtype_hint=tf.float32)
matrix = tf.convert_to_tensor(matrix, name='matrix', dtype=dtype)
if tensorshape_util.rank(matrix.shape) is None:
raise NotImplementedError(
'Rank of `matrix` must be known statically')
max_rank = tf.convert_to_tensor(max_rank,
name='max_rank',
dtype=tf.int64)
max_rank = tf.minimum(
max_rank,
prefer_static.shape(matrix, out_type=tf.int64)[-1])
diag_rtol = tf.convert_to_tensor(diag_rtol,
dtype=dtype,
name='diag_rtol')
matrix_diag = tf.linalg.diag_part(matrix)
# matrix is P.D., therefore all matrix_diag > 0, so we don't need abs.
orig_error = tf.reduce_max(matrix_diag, axis=-1)
def cond(m, pchol, perm, matrix_diag):
"""Condition for `tf.while_loop` continuation."""
del pchol
del perm
error = tf.linalg.norm(matrix_diag, ord=1, axis=-1)
max_err = tf.reduce_max(error / orig_error)
return (m < max_rank) & (tf.equal(m, 0) | (max_err > diag_rtol))
batch_dims = tensorshape_util.rank(matrix.shape) - 2
def batch_gather(params, indices, axis=-1):
return tf.gather(params, indices, axis=axis, batch_dims=batch_dims)
def body(m, pchol, perm, matrix_diag):
"""Body of a single `tf.while_loop` iteration."""
# Here is roughly a numpy, non-batched version of what's going to happen.
# (See also Algorithm 1 of Harbrecht et al.)
# 1: maxi = np.argmax(matrix_diag[perm[m:]]) + m
# 2: maxval = matrix_diag[perm][maxi]
# 3: perm[m], perm[maxi] = perm[maxi], perm[m]
# 4: row = matrix[perm[m]][perm[m + 1:]]
# 5: row -= np.sum(pchol[:m][perm[m + 1:]] * pchol[:m][perm[m]]], axis=-2)
# 6: pivot = np.sqrt(maxval); row /= pivot
# 7: row = np.concatenate([[[pivot]], row], -1)
# 8: matrix_diag[perm[m:]] -= row**2
# 9: pchol[m, perm[m:]] = row
# Find the maximal position of the (remaining) permuted diagonal.
# Steps 1, 2 above.
permuted_diag = batch_gather(matrix_diag, perm[..., m:])
maxi = tf.argmax(permuted_diag, axis=-1,
output_type=tf.int64)[..., tf.newaxis]
maxval = batch_gather(permuted_diag, maxi)
maxi = maxi + m
maxval = maxval[..., 0]
# Update perm: Swap perm[...,m] with perm[...,maxi]. Step 3 above.
perm = _swap_m_with_i(perm, m, maxi)
# Step 4.
row = batch_gather(matrix, perm[..., m:m + 1], axis=-2)
row = batch_gather(row, perm[..., m + 1:])
# Step 5.
prev_rows = pchol[..., :m, :]
prev_rows_perm_m_onward = batch_gather(prev_rows, perm[...,
m + 1:])
prev_rows_pivot_col = batch_gather(prev_rows, perm[..., m:m + 1])
row -= tf.reduce_sum(prev_rows_perm_m_onward * prev_rows_pivot_col,
axis=-2)[..., tf.newaxis, :]
# Step 6.
pivot = tf.sqrt(maxval)[..., tf.newaxis, tf.newaxis]
# Step 7.
row = tf.concat([pivot, row / pivot], axis=-1)
# TODO(b/130899118): Pad grad fails with int64 paddings.
# Step 8.
paddings = tf.concat([
tf.zeros([prefer_static.rank(pchol) - 1, 2], dtype=tf.int32),
[[tf.cast(m, tf.int32), 0]]
],
axis=0)
diag_update = tf.pad(row**2, paddings=paddings)[..., 0, :]
reverse_perm = _invert_permutation(perm)
matrix_diag -= batch_gather(diag_update, reverse_perm)
# Step 9.
row = tf.pad(row, paddings=paddings)
# TODO(bjp): Defer the reverse permutation all-at-once at the end?
row = batch_gather(row, reverse_perm)
pchol_shape = pchol.shape
pchol = tf.concat([pchol[..., :m, :], row, pchol[..., m + 1:, :]],
axis=-2)
tensorshape_util.set_shape(pchol, pchol_shape)
return m + 1, pchol, perm, matrix_diag
m = np.int64(0)
pchol = tf.zeros_like(matrix[..., :max_rank, :])
matrix_shape = prefer_static.shape(matrix, out_type=tf.int64)
perm = tf.broadcast_to(prefer_static.range(matrix_shape[-1]),
matrix_shape[:-1])
_, pchol, _, _ = tf.while_loop(cond=cond,
body=body,
loop_vars=(m, pchol, perm, matrix_diag))
pchol = tf.linalg.matrix_transpose(pchol)
tensorshape_util.set_shape(
pchol, tensorshape_util.concatenate(matrix_diag.shape, [None]))
return pchol
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
def __init__(self,
temperature,
logits=None,
probs=None,
validate_args=False,
allow_nan_stats=True,
name='RelaxedBernoulli'):
"""Construct RelaxedBernoulli distributions.
Args:
temperature: An 0-D `Tensor`, representing the temperature
of a set of RelaxedBernoulli distributions. The temperature should be
positive.
logits: An N-D `Tensor` representing the log-odds
of a positive event. Each entry in the `Tensor` parametrizes
an independent RelaxedBernoulli distribution where the probability of an
event is sigmoid(logits). Only one of `logits` or `probs` should be
passed in.
probs: An N-D `Tensor` representing the probability of a positive event.
Each entry in the `Tensor` parameterizes an independent Bernoulli
distribution. Only one of `logits` or `probs` should be passed in.
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 both `probs` and `logits` are passed, or if neither.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([logits, probs, temperature],
tf.float32)
self._temperature = tf.convert_to_tensor(temperature,
name='temperature',
dtype=dtype)
if validate_args:
with tf.control_dependencies(
[assert_util.assert_positive(temperature)]):
self._temperature = tf.identity(self._temperature)
self._logits, self._probs = distribution_util.get_logits_and_probs(
logits=logits,
probs=probs,
validate_args=validate_args,
dtype=dtype)
super(RelaxedBernoulli, self).__init__(
distribution=logistic.Logistic(self._logits /
self._temperature,
1. / self._temperature,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
name=name + '/Logistic'),
bijector=sigmoid_bijector.Sigmoid(validate_args=validate_args),
validate_args=validate_args,
name=name)
self._parameters = parameters
def __init__(self,
loc=None,
scale_diag=None,
scale_identity_multiplier=None,
validate_args=False,
allow_nan_stats=True,
name="VectorLaplaceDiag"):
"""Construct Vector Laplace 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 = 2 * scale @ scale.T`.
```none
scale = diag(scale_diag + scale_identity_multiplier * ones(k))
```
where:
* `scale_diag.shape = [k]`, and,
* `scale_identity_multiplier.shape = []`.
Additional leading dimensions (if any) will index batches.
If both `scale_diag` and `scale_identity_multiplier` are `None`, then
`scale` is the Identity matrix.
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_diag: Non-zero, floating-point `Tensor` representing a diagonal
matrix added to `scale`. May have shape `[B1, ..., Bb, k]`, `b >= 0`,
and characterizes `b`-batches of `k x k` diagonal matrices added to
`scale`. When both `scale_identity_multiplier` and `scale_diag` are
`None` then `scale` is the `Identity`.
scale_identity_multiplier: Non-zero, floating-point `Tensor` representing
a scaled-identity-matrix added to `scale`. May have shape
`[B1, ..., Bb]`, `b >= 0`, and characterizes `b`-batches of scaled
`k x k` identity matrices added to `scale`. When both
`scale_identity_multiplier` and `scale_diag` are `None` then `scale` is
the `Identity`.
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 at most `scale_identity_multiplier` is specified.
"""
parameters = dict(locals())
with tf.name_scope(name):
with tf.name_scope("init"):
dtype = dtype_util.common_dtype(
[loc, scale_diag, scale_identity_multiplier], tf.float32)
# No need to validate_args while making diag_scale. The returned
# LinearOperatorDiag has an assert_non_singular method that is called by
# the Bijector.
scale = distribution_util.make_diag_scale(
loc=loc,
scale_diag=scale_diag,
scale_identity_multiplier=scale_identity_multiplier,
validate_args=False,
assert_positive=False,
dtype=dtype)
super(VectorLaplaceDiag,
self).__init__(loc=loc,
scale=scale,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
name=name)
self._parameters = parameters
def __init__(self,
shift=None,
scale_identity_multiplier=None,
scale_diag=None,
scale_tril=None,
scale_perturb_factor=None,
scale_perturb_diag=None,
adjoint=False,
validate_args=False,
name="affine",
dtype=None):
"""Instantiates the `Affine` bijector.
This `Bijector` is initialized with `shift` `Tensor` and `scale` arguments,
giving the forward operation:
```none
Y = g(X) = scale @ X + shift
```
where the `scale` term is logically equivalent to:
```python
scale = (
scale_identity_multiplier * tf.diag(tf.ones(d)) +
tf.diag(scale_diag) +
scale_tril +
scale_perturb_factor @ diag(scale_perturb_diag) @
tf.transpose([scale_perturb_factor])
)
```
If none of `scale_identity_multiplier`, `scale_diag`, or `scale_tril` are
specified then `scale += IdentityMatrix`. Otherwise specifying a
`scale` argument has the semantics of `scale += Expand(arg)`, i.e.,
`scale_diag != None` means `scale += tf.diag(scale_diag)`.
Args:
shift: Floating-point `Tensor`. If this is set to `None`, no shift is
applied.
scale_identity_multiplier: floating point rank 0 `Tensor` representing a
scaling done to the identity matrix.
When `scale_identity_multiplier = scale_diag = scale_tril = None` then
`scale += IdentityMatrix`. Otherwise no scaled-identity-matrix is added
to `scale`.
scale_diag: Floating-point `Tensor` representing the diagonal matrix.
`scale_diag` has shape `[N1, N2, ... k]`, which represents a k x k
diagonal matrix.
When `None` no diagonal term is added to `scale`.
scale_tril: Floating-point `Tensor` representing the lower triangular
matrix. `scale_tril` has shape `[N1, N2, ... k, k]`, which represents a
k x k lower triangular matrix.
When `None` no `scale_tril` term is added to `scale`.
The upper triangular elements above the diagonal are ignored.
scale_perturb_factor: Floating-point `Tensor` representing factor matrix
with last two dimensions of shape `(k, r)`. When `None`, no rank-r
update is added to `scale`.
scale_perturb_diag: Floating-point `Tensor` representing the diagonal
matrix. `scale_perturb_diag` has shape `[N1, N2, ... r]`, which
represents an `r x r` diagonal matrix. When `None` low rank updates will
take the form `scale_perturb_factor * scale_perturb_factor.T`.
adjoint: Python `bool` indicating whether to use the `scale` matrix as
specified or its adjoint.
Default value: `False`.
validate_args: Python `bool` indicating whether arguments should be
checked for correctness.
name: Python `str` name given to ops managed by this object.
dtype: `tf.DType` to prefer when converting args to `Tensor`s. Else, we
fall back to a common dtype inferred from the args, finally falling back
to float32.
Raises:
ValueError: if `perturb_diag` is specified but not `perturb_factor`.
TypeError: if `shift` has different `dtype` from `scale` arguments.
"""
# Ambiguous definition of low rank update.
if scale_perturb_diag is not None and scale_perturb_factor is None:
raise ValueError("When scale_perturb_diag is specified, "
"scale_perturb_factor must be specified.")
# Special case, only handling a scaled identity matrix. We don't know its
# dimensions, so this is special cased.
# We don't check identity_multiplier, since below we set it to 1. if all
# other scale args are None.
self._is_only_identity_multiplier = (scale_tril is None
and scale_diag is None
and scale_perturb_factor is None)
with tf.name_scope(name) as name:
self._name = name
self._validate_args = validate_args
if dtype is None:
dtype = dtype_util.common_dtype([
shift, scale_identity_multiplier, scale_diag, scale_tril,
scale_perturb_diag, scale_perturb_factor
], tf.float32)
if shift is not None:
shift = tf.convert_to_tensor(shift, name="shift", dtype=dtype)
self._shift = shift
# When no args are specified, pretend the scale matrix is the identity
# matrix.
if (self._is_only_identity_multiplier
and scale_identity_multiplier is None):
scale_identity_multiplier = tf.convert_to_tensor(1.,
dtype=dtype)
# self._create_scale_operator returns a LinearOperator in all cases
# except if self._is_only_identity_multiplier; in which case it
# returns a scalar Tensor.
scale = self._create_scale_operator(
identity_multiplier=scale_identity_multiplier,
diag=scale_diag,
tril=scale_tril,
perturb_diag=scale_perturb_diag,
perturb_factor=scale_perturb_factor,
shift=shift,
validate_args=validate_args,
dtype=dtype)
if (scale is not None and not self._is_only_identity_multiplier
and not dtype_util.SKIP_DTYPE_CHECKS):
if (shift is not None and
not dtype_util.base_equal(shift.dtype, scale.dtype)):
raise TypeError(
"shift.dtype ({}) is incompatible with scale.dtype ({})."
.format(shift.dtype, scale.dtype))
self._scale = scale
self._adjoint = adjoint
super(Affine, self).__init__(forward_min_event_ndims=1,
is_constant_jacobian=True,
dtype=dtype,
validate_args=validate_args,
name=name)
def __init__(self,
loc=None,
scale=None,
validate_args=False,
allow_nan_stats=True,
name='MultivariateNormalLinearOperator'):
"""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`.
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: Instance of `LinearOperator` with same `dtype` as `loc` and shape
`[B1, ..., Bb, k, k]`.
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.
Raises:
ValueError: if `scale` is unspecified.
TypeError: if not `scale.dtype.is_floating`
"""
parameters = dict(locals())
if scale is None:
raise ValueError('Missing required `scale` parameter.')
if not dtype_util.is_floating(scale.dtype):
raise TypeError(
'`scale` parameter must have floating-point dtype.')
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([loc, scale],
dtype_hint=tf.float32)
# Since expand_dims doesn't preserve constant-ness, we obtain the
# non-dynamic value if possible.
loc = tensor_util.convert_nonref_to_tensor(loc,
dtype=dtype,
name='loc')
batch_shape, event_shape = distribution_util.shapes_from_loc_and_scale(
loc, scale)
super(MultivariateNormalLinearOperator, self).__init__(
distribution=normal.Normal(loc=tf.zeros([], dtype=dtype),
scale=tf.ones([], dtype=dtype)),
bijector=affine_linear_operator_bijector.AffineLinearOperator(
shift=loc, scale=scale, validate_args=validate_args),
batch_shape=batch_shape,
event_shape=event_shape,
validate_args=validate_args,
name=name)
self._parameters = parameters
def __init__(self,
loc,
atol=None,
rtol=None,
is_vector=False,
validate_args=False,
allow_nan_stats=True,
parameters=None,
name="_BaseDeterministic"):
"""Initialize a batch of `_BaseDeterministic` distributions.
The `atol` and `rtol` parameters allow for some slack in `pmf`, `cdf`
computations, e.g. due to floating-point error.
```
pmf(x; loc)
= 1, if Abs(x - loc) <= atol + rtol * Abs(loc),
= 0, otherwise.
```
Args:
loc: Numeric `Tensor`. The point (or batch of points) on which this
distribution is supported.
atol: Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
shape. The absolute tolerance for comparing closeness to `loc`.
Default is `0`.
rtol: Non-negative `Tensor` of same `dtype` as `loc` and broadcastable
shape. The relative tolerance for comparing closeness to `loc`.
Default is `0`.
is_vector: Python `bool`. If `True`, this is for `VectorDeterministic`,
else `Deterministic`.
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.
parameters: Dict of locals to facilitate copy construction.
name: Python `str` name prefixed to Ops created by this class.
Raises:
ValueError: If `loc` is a scalar.
"""
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([loc, atol, rtol],
dtype_hint=tf.float32)
self._loc = tensor_util.convert_nonref_to_tensor(loc,
dtype_hint=dtype,
name="loc")
self._atol = tensor_util.convert_nonref_to_tensor(
0 if atol is None else atol, dtype=dtype, name="atol")
self._rtol = tensor_util.convert_nonref_to_tensor(
0 if rtol is None else rtol, dtype=dtype, name="rtol")
self._is_vector = is_vector
super(_BaseDeterministic,
self).__init__(dtype=self._loc.dtype,
reparameterization_type=(
reparameterization.FULLY_REPARAMETERIZED
if self._loc.dtype.is_floating else
reparameterization.NOT_REPARAMETERIZED),
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name)
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)
loc = tf.convert_to_tensor(loc, name="loc", dtype=dtype)
scale = tf.convert_to_tensor(scale, name="scale", dtype=dtype)
low = tf.convert_to_tensor(low, name="low", dtype=dtype)
high = tf.convert_to_tensor(high, name="high", dtype=dtype)
dtype_util.assert_same_float_dtype([loc, scale, low, high])
self._broadcast_batch_shape = distribution_util.get_broadcast_shape(
loc, scale, low, high)
# Broadcast all parameters to the same shape
broadcast_ones = tf.ones(shape=self._broadcast_batch_shape,
dtype=scale.dtype)
self._scale = scale * broadcast_ones
self._loc = loc * broadcast_ones
self._low = low * broadcast_ones
self._high = high * broadcast_ones
with tf.control_dependencies(
[self._validate()] if validate_args else []):
self._loc = tf.identity(self._loc)
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,
shift=None,
scale=None,
log_scale=None,
validate_args=False,
name='affine_scalar'):
"""Instantiates the `AffineScalar` bijector.
This `Bijector` is initialized with `shift` `Tensor` and `scale` arguments,
giving the forward operation:
```none
Y = g(X) = scale * X + shift
```
Alternatively, you can specify `log_scale` instead of `scale` for slighly
better numerics with tiny scales. Note that when using `log_scale` it is
currently impossible to specify a negative scale.
If `scale` or `log_scale` are not specified, then the bijector has the
semantics of `scale = 1.`. Similarly, if `shift` is not specified, then the
bijector has the semantics of `shift = 0.`.
Args:
shift: Floating-point `Tensor`. If this is set to `None`, no shift is
applied.
scale: Floating-point `Tensor`. If this is set to `None`, no scale is
applied. This should not be set if `log_scale` is set.
log_scale: Floating-point `Tensor`. Logarithm of the scale. If this is set
to `None`, no scale is applied. This should not be set if `scale` is
set.
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 `scale` and `log_scale` are specified.
"""
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([shift, scale, log_scale],
dtype_hint=tf.float32)
if scale is not None and log_scale is not None:
raise ValueError(
'At most one of `scale` and `log_scale` should be '
'specified')
self._shift = tensor_util.convert_nonref_to_tensor(shift,
dtype=dtype,
name='shift')
self._scale = tensor_util.convert_nonref_to_tensor(scale,
dtype=dtype,
name='scale')
self._log_scale = tensor_util.convert_nonref_to_tensor(
log_scale, dtype=dtype, name='log_scale')
super(AffineScalar, self).__init__(forward_min_event_ndims=0,
is_constant_jacobian=True,
validate_args=validate_args,
dtype=dtype,
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)
mean_direction = tf.convert_to_tensor(mean_direction,
name='mean_direction',
dtype=dtype)
concentration = tf.convert_to_tensor(concentration,
name='concentration',
dtype=dtype)
assertions = [
assert_util.assert_non_negative(
concentration,
message='`concentration` must be non-negative'),
assert_util.assert_greater(
tf.shape(mean_direction)[-1],
1,
message='`mean_direction` may not have scalar event shape'
),
assert_util.assert_near(
1.,
tf.linalg.norm(mean_direction, axis=-1),
message='`mean_direction` must be unit-length')
] if validate_args else []
static_event_dim = tf.compat.dimension_value(
tensorshape_util.with_rank_at_least(mean_direction.shape,
1)[-1])
if static_event_dim is not None and static_event_dim > 5:
raise ValueError('vMF ndims > 5 is not currently supported')
elif validate_args:
assertions += [
assert_util.assert_less_equal(
tf.shape(mean_direction)[-1],
5,
message='vMF ndims > 5 is not currently supported')
]
with tf.control_dependencies(assertions):
self._mean_direction = tf.identity(mean_direction)
self._concentration = tf.identity(concentration)
dtype_util.assert_same_float_dtype(
[self._mean_direction, self._concentration])
# 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,
loc=None,
covariance_matrix=None,
validate_args=False,
allow_nan_stats=True,
name="MultivariateNormalFullCovariance"):
"""Construct Multivariate Normal distribution on `R^k`.
The `batch_shape` is the broadcast shape between `loc` and
`covariance_matrix` arguments.
The `event_shape` is given by last dimension of the matrix implied by
`covariance_matrix`. The last dimension of `loc` (if provided) must
broadcast with this.
A non-batch `covariance_matrix` matrix is a `k x k` symmetric positive
definite matrix. In other words it is (real) symmetric with all eigenvalues
strictly positive.
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.
covariance_matrix: Floating-point, symmetric positive definite `Tensor` of
same `dtype` as `loc`. The strict upper triangle of `covariance_matrix`
is ignored, so if `covariance_matrix` is not symmetric no error will be
raised (unless `validate_args is True`). `covariance_matrix` 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 `covariance_matrix` are specified.
"""
parameters = dict(locals())
# Convert the covariance_matrix up to a scale_tril and call MVNTriL.
with tf.name_scope(name) as name:
with tf.name_scope("init"):
dtype = dtype_util.common_dtype([loc, covariance_matrix],
tf.float32)
loc = loc if loc is None else tf.convert_to_tensor(
loc, name="loc", dtype=dtype)
if covariance_matrix is None:
scale_tril = None
else:
covariance_matrix = tf.convert_to_tensor(
covariance_matrix,
name="covariance_matrix",
dtype=dtype)
if validate_args:
covariance_matrix = distribution_util.with_dependencies(
[
assert_util.assert_near(
covariance_matrix,
tf.linalg.matrix_transpose(
covariance_matrix),
message="Matrix was not symmetric")
], covariance_matrix)
# No need to validate that covariance_matrix is non-singular.
# LinearOperatorLowerTriangular has an assert_non_singular method that
# is called by the Bijector.
# However, cholesky() ignores the upper triangular part, so we do need
# to separately assert symmetric.
scale_tril = tf.linalg.cholesky(covariance_matrix)
super(MultivariateNormalFullCovariance,
self).__init__(loc=loc,
scale_tril=scale_tril,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
name=name)
self._parameters = parameters
def __init__(self,
df,
loc=None,
scale_identity_multiplier=None,
scale_diag=None,
scale_tril=None,
scale_perturb_factor=None,
scale_perturb_diag=None,
validate_args=False,
allow_nan_stats=True,
name="VectorStudentT"):
"""Instantiates the vector Student's t-distributions on `R^k`.
The `batch_shape` is the broadcast between `df.batch_shape` and
`Affine.batch_shape` where `Affine` is constructed from `loc` and
`scale_*` arguments.
The `event_shape` is the event shape of `Affine.event_shape`.
Args:
df: Floating-point `Tensor`. The degrees of freedom of the
distribution(s). `df` must contain only positive values. Must be
scalar if `loc`, `scale_*` imply non-scalar batch_shape or must have the
same `batch_shape` implied by `loc`, `scale_*`.
loc: Floating-point `Tensor`. If this is set to `None`, no `loc` is
applied.
scale_identity_multiplier: floating point rank 0 `Tensor` representing a
scaling done to the identity matrix. When `scale_identity_multiplier =
scale_diag=scale_tril = None` then `scale += IdentityMatrix`. Otherwise
no scaled-identity-matrix is added to `scale`.
scale_diag: Floating-point `Tensor` representing the diagonal matrix.
`scale_diag` has shape [N1, N2, ..., k], which represents a k x k
diagonal matrix. When `None` no diagonal term is added to `scale`.
scale_tril: Floating-point `Tensor` representing the diagonal matrix.
`scale_diag` has shape [N1, N2, ..., k, k], which represents a k x k
lower triangular matrix. When `None` no `scale_tril` term is added to
`scale`. The upper triangular elements above the diagonal are ignored.
scale_perturb_factor: Floating-point `Tensor` representing factor matrix
with last two dimensions of shape `(k, r)`. When `None`, no rank-r
update is added to `scale`.
scale_perturb_diag: Floating-point `Tensor` representing the diagonal
matrix. `scale_perturb_diag` has shape [N1, N2, ..., r], which
represents an r x r Diagonal matrix. When `None` low rank updates will
take the form `scale_perturb_factor * scale_perturb_factor.T`.
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())
args = [
df, loc, scale_identity_multiplier, scale_diag, scale_tril,
scale_perturb_factor, scale_perturb_diag
]
with tf.name_scope(name) as name:
with tf.name_scope("init"):
dtype = dtype_util.common_dtype(args, tf.float32)
df = tf.convert_to_tensor(df, name="df", dtype=dtype)
# The shape of the _VectorStudentT distribution is governed by the
# relationship between df.batch_shape and affine.batch_shape. In
# pseudocode the basic procedure is:
# if df.batch_shape is scalar:
# if affine.batch_shape is not scalar:
# # broadcast distribution.sample so
# # it has affine.batch_shape.
# self.batch_shape = affine.batch_shape
# else:
# if affine.batch_shape is scalar:
# # let affine broadcasting do its thing.
# self.batch_shape = df.batch_shape
# All of the above magic is actually handled by TransformedDistribution.
# Here we really only need to collect the affine.batch_shape and decide
# what we're going to pass in to TransformedDistribution's
# (override) batch_shape arg.
affine = affine_bijector.Affine(
shift=loc,
scale_identity_multiplier=scale_identity_multiplier,
scale_diag=scale_diag,
scale_tril=scale_tril,
scale_perturb_factor=scale_perturb_factor,
scale_perturb_diag=scale_perturb_diag,
validate_args=validate_args,
dtype=dtype)
distribution = student_t.StudentT(
df=df,
loc=tf.zeros([], dtype=affine.dtype),
scale=tf.ones([], dtype=affine.dtype))
batch_shape, override_event_shape = (
distribution_util.shapes_from_loc_and_scale(
affine.shift, affine.scale))
override_batch_shape = distribution_util.pick_vector(
distribution.is_scalar_batch(), batch_shape,
tf.constant([], dtype=tf.int32))
super(_VectorStudentT,
self).__init__(distribution=distribution,
bijector=affine,
batch_shape=override_batch_shape,
event_shape=override_event_shape,
validate_args=validate_args,
name=name)
self._parameters = parameters
def __init__(self,
loc=None,
scale_diag=None,
scale_identity_multiplier=None,
skewness=None,
tailweight=None,
distribution=None,
validate_args=False,
allow_nan_stats=True,
name="VectorSinhArcsinhDiag"):
"""Construct VectorSinhArcsinhDiag distribution on `R^k`.
The arguments `scale_diag` and `scale_identity_multiplier` combine to
define the diagonal `scale` referred to in this class docstring:
```none
scale = diag(scale_diag + scale_identity_multiplier * ones(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
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_diag: Non-zero, floating-point `Tensor` representing a diagonal
matrix added to `scale`. May have shape `[B1, ..., Bb, k]`, `b >= 0`,
and characterizes `b`-batches of `k x k` diagonal matrices added to
`scale`. When both `scale_identity_multiplier` and `scale_diag` are
`None` then `scale` is the `Identity`.
scale_identity_multiplier: Non-zero, floating-point `Tensor` representing
a scale-identity-matrix added to `scale`. May have shape
`[B1, ..., Bb]`, `b >= 0`, and characterizes `b`-batches of scale
`k x k` identity matrices added to `scale`. When both
`scale_identity_multiplier` and `scale_diag` are `None` then `scale`
is the `Identity`.
skewness: Skewness parameter. floating-point `Tensor` with shape
broadcastable with `event_shape`.
tailweight: Tailweight parameter. floating-point `Tensor` with shape
broadcastable with `event_shape`.
distribution: `tf.Distribution`-like instance. Distribution from which `k`
iid samples are used as input to transformation `F`. Default is
`tfd.Normal(loc=0., scale=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 VectorSinhArcsinhDiag 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.
Raises:
ValueError: if at most `scale_identity_multiplier` is specified.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([
loc, scale_diag, scale_identity_multiplier, skewness,
tailweight
], tf.float32)
loc = loc if loc is None else tf.convert_to_tensor(
loc, name="loc", dtype=dtype)
tailweight = 1. if tailweight is None else tailweight
skewness = 0. if skewness is None else skewness
# Recall, with Z a random variable,
# Y := loc + C * F(Z),
# F(Z) := Sinh( (Arcsinh(Z) + skewness) * tailweight )
# F_0(Z) := Sinh( Arcsinh(Z) * tailweight )
# C := 2 * scale / F_0(2)
# Construct shapes and 'scale' out of the scale_* and loc kwargs.
# scale_linop is only an intermediary to:
# 1. get shapes from looking at loc and the two scale args.
# 2. combine scale_diag with scale_identity_multiplier, which gives us
# 'scale', which in turn gives us 'C'.
scale_linop = distribution_util.make_diag_scale(
loc=loc,
scale_diag=scale_diag,
scale_identity_multiplier=scale_identity_multiplier,
validate_args=False,
assert_positive=False,
dtype=dtype)
batch_shape, event_shape = distribution_util.shapes_from_loc_and_scale(
loc, scale_linop)
# scale_linop.diag_part() is efficient since it is a diag type linop.
scale_diag_part = scale_linop.diag_part()
dtype = scale_diag_part.dtype
if distribution is None:
distribution = normal.Normal(loc=tf.zeros([], dtype=dtype),
scale=tf.ones([], dtype=dtype),
allow_nan_stats=allow_nan_stats)
else:
asserts = distribution_util.maybe_check_scalar_distribution(
distribution, dtype, validate_args)
if asserts:
scale_diag_part = distribution_util.with_dependencies(
asserts, scale_diag_part)
# Make the SAS bijector, 'F'.
skewness = tf.convert_to_tensor(skewness,
dtype=dtype,
name="skewness")
tailweight = tf.convert_to_tensor(tailweight,
dtype=dtype,
name="tailweight")
f = sinh_arcsinh_bijector.SinhArcsinh(skewness=skewness,
tailweight=tailweight)
affine = affine_bijector.Affine(shift=loc,
scale_diag=scale_diag_part,
validate_args=validate_args)
bijector = chain_bijector.Chain([affine, f])
super(VectorSinhArcsinhDiag,
self).__init__(distribution=distribution,
bijector=bijector,
batch_shape=batch_shape,
event_shape=event_shape,
validate_args=validate_args,
name=name)
self._parameters = parameters
self._loc = loc
self._scale = scale_linop
self._tailweight = tailweight
self._skewness = skewness
