def _multi_gamma_sequence(self, a, p, name="multi_gamma_sequence"):
"""Creates sequence used in multivariate (di)gamma; shape = shape(a)+[p]."""
with self._name_and_control_scope(name):
# Linspace only takes scalars, so we'll add in the offset afterwards.
seq = tf.linspace(tf.constant(0., dtype=self.dtype), 0.5 - 0.5 * p,
tf.cast(p, tf.int32))
return seq + tf.expand_dims(a, [-1])
def _mean(self):
concentration = tf.convert_to_tensor(self.concentration)
lim = tf.ones([], dtype=self.dtype)
valid = concentration < lim
safe_conc = tf.where(valid, concentration, tf.constant(.5, self.dtype))
result = lambda: self.loc + self.scale / (1 - safe_conc)
if self.allow_nan_stats:
return tf.where(valid, result(),
tf.constant(float('nan'), self.dtype))
with tf.control_dependencies([
assert_util.assert_less(
concentration,
lim,
message='`mean` is undefined when `concentration >= 1`')
]):
return result()
def maybe_assert_bernoulli_param_correctness(
is_init, validate_args, probs, logits):
"""Return assertions for `Bernoulli`-type distributions."""
if is_init:
x, name = (probs, 'probs') if logits is None else (logits, 'logits')
if not dtype_util.is_floating(x.dtype):
raise TypeError(
'Argument `{}` must having floating type.'.format(name))
if not validate_args:
return []
assertions = []
if probs is not None:
if is_init != tensor_util.is_ref(probs):
probs = tf.convert_to_tensor(probs)
one = tf.constant(1., probs.dtype)
assertions += [
assert_util.assert_non_negative(
probs, message='probs has components less than 0.'),
assert_util.assert_less_equal(
probs, one, message='probs has components greater than 1.')
]
return assertions
def _maybe_assert_valid_y(self, y):
if not self.validate_args:
return []
is_positive = assert_util.assert_non_negative(
y, message='Inverse transformation input must be greater than 0.')
less_than_one = assert_util.assert_less_equal(
y,
tf.constant(1., y.dtype),
message='Inverse transformation input must be less than or equal to 1.')
return [is_positive, less_than_one]
def _variance(self):
concentration = tf.convert_to_tensor(self.concentration)
lim = tf.constant(.5, self.dtype)
valid = concentration < lim
safe_conc = tf.where(valid, concentration,
tf.constant(.25, self.dtype))
result = lambda: self.scale**2 / ((1 - safe_conc)**2 *
(1 - 2 * safe_conc))
if self.allow_nan_stats:
return tf.where(valid, result(),
tf.constant(float('nan'), self.dtype))
with tf.control_dependencies([
assert_util.assert_less(
concentration,
lim,
message=
'`variance` is undefined when `concentration >= 0.5`')
]):
return result()
def _forward_log_det_jacobian(self, x):
# is_constant_jacobian = True for this bijector, hence the
# `log_det_jacobian` need only be specified for a single input, as this will
# be tiled to match `event_ndims`.
if self.scale is None:
return tf.constant(0., dtype=dtype_util.base_dtype(x.dtype))
with tf.control_dependencies(self._maybe_collect_assertions() if self.
validate_args else []):
return self.scale.log_abs_determinant()
def _log_prob(self, x):
scale = tf.convert_to_tensor(self.scale)
concentration = tf.convert_to_tensor(self.concentration)
z = self._z(x, scale, concentration)
eq_zero = tf.equal(concentration,
0) # Concentration = 0 ==> Exponential.
nonzero_conc = tf.where(eq_zero, tf.constant(1, self.dtype),
concentration)
where_nonzero = (1 / nonzero_conc + 1) * tf.math.log1p(
nonzero_conc * z)
return -tf.math.log(scale) - tf.where(eq_zero, z, where_nonzero)
def _ndtr(x):
"""Implements ndtr core logic."""
half_sqrt_2 = tf.constant(
0.5 * np.sqrt(2.), dtype=x.dtype, name="half_sqrt_2")
w = x * half_sqrt_2
z = tf.abs(w)
y = tf.where(
z < half_sqrt_2,
1. + tf.math.erf(w),
tf.where(w > 0., 2. - tf.math.erfc(z), tf.math.erfc(z)))
return 0.5 * y
def _log_cdf(self, x):
scale = tf.convert_to_tensor(self.scale)
concentration = tf.convert_to_tensor(self.concentration)
z = self._z(x, scale, concentration)
eq_zero = tf.equal(concentration,
0) # Concentration = 0 ==> Exponential.
nonzero_conc = tf.where(eq_zero, tf.constant(1, self.dtype),
concentration)
where_nonzero = tf.math.log1p(-(1 + nonzero_conc * z)**(-1 /
nonzero_conc))
where_zero = tf.math.log1p(-tf.exp(-z))
return tf.where(eq_zero, where_zero, where_nonzero)
def _variance(self):
concentration = tf.convert_to_tensor(self.concentration)
scale = tf.convert_to_tensor(self.scale)
var = (tf.square(scale) / tf.square(concentration - 1.) /
(concentration - 2.))
if self.allow_nan_stats:
assertions = []
else:
assertions = [
assert_util.assert_less(
tf.constant(2., dtype=self.dtype),
concentration,
message='variance undefined when any concentration <= 2')
]
with tf.control_dependencies(assertions):
return tf.where(concentration > 2., var,
dtype_util.as_numpy_dtype(self.dtype)(np.nan))
def _sample_n(self, n, seed=None):
# Inversion samples via inverse CDF.
loc = tf.convert_to_tensor(self.loc)
scale = tf.convert_to_tensor(self.scale)
concentration = tf.convert_to_tensor(self.concentration)
# Pre-broadcast to ensure we draw enough randomness.
sample_shp = tf.concat(
[[n],
self._batch_shape_tensor(
loc=loc, scale=scale, concentration=concentration)],
axis=0)
logu = tf.math.log1p(
-tf.random.uniform(sample_shp, dtype=self.dtype, seed=seed))
eq_zero = tf.equal(concentration, 0)
safe_conc = tf.where(eq_zero, tf.constant(1, dtype=self.dtype),
concentration)
where_nonzero = loc + scale / safe_conc * tf.math.expm1(
-safe_conc * logu)
where_zero = loc - scale * logu
return tf.where(eq_zero, where_zero, where_nonzero)
def _forward(self, x):
x = tf.convert_to_tensor(x, name='x')
batch_shape = prefer_static.shape(x)[:-1]
# Pad zeros on the top row and right column.
y = fill_triangular.FillTriangular().forward(x)
rank = prefer_static.rank(y)
paddings = tf.concat([
tf.zeros(shape=(rank - 2, 2), dtype=tf.int32),
tf.constant([[1, 0], [0, 1]], dtype=tf.int32)
],
axis=0)
y = tf.pad(y, paddings)
# Set diagonal to 1s.
n = prefer_static.shape(y)[-1]
diag = tf.ones(tf.concat([batch_shape, [n]], axis=-1), dtype=x.dtype)
y = tf.linalg.set_diag(y, diag)
# Normalize each row to have Euclidean (L2) norm 1.
y /= tf.norm(y, axis=-1)[..., tf.newaxis]
return y
def maybe_assert_negative_binomial_param_correctness(is_init, validate_args,
total_count, probs,
logits):
"""Return assertions for `NegativeBinomial`-type distributions."""
if is_init:
x, name = (probs, 'probs') if logits is None else (logits, 'logits')
if not dtype_util.is_floating(x.dtype):
raise TypeError(
'Argument `{}` must having floating type.'.format(name))
if not validate_args:
return []
assertions = []
if is_init != tensor_util.is_ref(total_count):
total_count = tf.convert_to_tensor(total_count)
assertions.extend([
assert_util.assert_non_negative(
total_count,
message='`total_count` has components less than 0.'),
distribution_util.assert_integer_form(
total_count,
message='`total_count` has fractional components.')
])
if probs is not None:
if is_init != tensor_util.is_ref(probs):
probs = tf.convert_to_tensor(probs)
one = tf.constant(1., probs.dtype)
assertions.extend([
assert_util.assert_non_negative(
probs, message='`probs` has components less than 0.'),
assert_util.assert_less_equal(
probs,
one,
message='`probs` has components greater than 1.')
])
return assertions
def _ndtri(p):
"""Implements ndtri core logic."""
# Constants used in piece-wise rational approximations. Taken from the cephes
# library:
# https://root.cern.ch/doc/v608/SpecFuncCephesInv_8cxx_source.html
p0 = list(reversed([-5.99633501014107895267E1,
9.80010754185999661536E1,
-5.66762857469070293439E1,
1.39312609387279679503E1,
-1.23916583867381258016E0]))
q0 = list(reversed([1.0,
1.95448858338141759834E0,
4.67627912898881538453E0,
8.63602421390890590575E1,
-2.25462687854119370527E2,
2.00260212380060660359E2,
-8.20372256168333339912E1,
1.59056225126211695515E1,
-1.18331621121330003142E0]))
p1 = list(reversed([4.05544892305962419923E0,
3.15251094599893866154E1,
5.71628192246421288162E1,
4.40805073893200834700E1,
1.46849561928858024014E1,
2.18663306850790267539E0,
-1.40256079171354495875E-1,
-3.50424626827848203418E-2,
-8.57456785154685413611E-4]))
q1 = list(reversed([1.0,
1.57799883256466749731E1,
4.53907635128879210584E1,
4.13172038254672030440E1,
1.50425385692907503408E1,
2.50464946208309415979E0,
-1.42182922854787788574E-1,
-3.80806407691578277194E-2,
-9.33259480895457427372E-4]))
p2 = list(reversed([3.23774891776946035970E0,
6.91522889068984211695E0,
3.93881025292474443415E0,
1.33303460815807542389E0,
2.01485389549179081538E-1,
1.23716634817820021358E-2,
3.01581553508235416007E-4,
2.65806974686737550832E-6,
6.23974539184983293730E-9]))
q2 = list(reversed([1.0,
6.02427039364742014255E0,
3.67983563856160859403E0,
1.37702099489081330271E0,
2.16236993594496635890E-1,
1.34204006088543189037E-2,
3.28014464682127739104E-4,
2.89247864745380683936E-6,
6.79019408009981274425E-9]))
def _create_polynomial(var, coeffs):
"""Compute n_th order polynomial via Horner's method."""
coeffs = np.array(coeffs, dtype_util.as_numpy_dtype(var.dtype))
if not coeffs.size:
return tf.zeros_like(var)
return coeffs[0] + _create_polynomial(var, coeffs[1:]) * var
maybe_complement_p = tf.where(p > -np.expm1(-2.), 1. - p, p)
# Write in an arbitrary value in place of 0 for p since 0 will cause NaNs
# later on. The result from the computation when p == 0 is not used so any
# number that doesn't result in NaNs is fine.
sanitized_mcp = tf.where(
maybe_complement_p <= 0.,
dtype_util.as_numpy_dtype(p.dtype)(0.5), maybe_complement_p)
# Compute x for p > exp(-2): x/sqrt(2pi) = w + w**3 P0(w**2)/Q0(w**2).
w = sanitized_mcp - 0.5
ww = w ** 2
x_for_big_p = w + w * ww * (_create_polynomial(ww, p0)
/ _create_polynomial(ww, q0))
x_for_big_p *= -np.sqrt(2. * np.pi)
# Compute x for p <= exp(-2): x = z - log(z)/z - (1/z) P(1/z) / Q(1/z),
# where z = sqrt(-2. * log(p)), and P/Q are chosen between two different
# arrays based on whether p < exp(-32).
z = tf.sqrt(-2. * tf.math.log(sanitized_mcp))
first_term = z - tf.math.log(z) / z
second_term_small_p = (
_create_polynomial(1. / z, p2) /
_create_polynomial(1. / z, q2) / z)
second_term_otherwise = (
_create_polynomial(1. / z, p1) /
_create_polynomial(1. / z, q1) / z)
x_for_small_p = first_term - second_term_small_p
x_otherwise = first_term - second_term_otherwise
x = tf.where(sanitized_mcp > np.exp(-2.), x_for_big_p,
tf.where(z >= 8.0, x_for_small_p, x_otherwise))
x = tf.where(p > 1. - np.exp(-2.), x, -x)
infinity_scalar = tf.constant(np.inf, dtype=p.dtype)
x_nan_replaced = tf.where(p <= 0.0, -infinity_scalar,
tf.where(p >= 1.0, infinity_scalar, x))
return x_nan_replaced
def _event_shape_tensor(self, loc=None):
del loc
return tf.constant([], dtype=tf.int32)
def _inverse_log_det_jacobian(self, y):
return tf.constant(0., dtype=dtype_util.base_dtype(y.dtype))
def _forward_log_det_jacobian(self, x):
return tf.constant(0., dtype=dtype_util.base_dtype(x.dtype))
def _inverse_log_det_jacobian(self, y):
# is_constant_jacobian = True for this bijector, hence the
# `log_det_jacobian` need only be specified for a single input, as this will
# be tiled to match `event_ndims`.
return tf.constant(0., dtype=dtype_util.base_dtype(y.dtype))
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`.
"""
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._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 will keep track of a static and dynamic version of
# self._is_{batch,event}_override. This way we can do more prior to graph
# execution, including possibly raising Python exceptions.
self._override_batch_shape = self._maybe_validate_shape_override(
batch_shape, distribution.is_scalar_batch(), validate_args,
"batch_shape")
self._is_batch_override = prefer_static.logical_not(
prefer_static.equal(
prefer_static.rank_from_shape(self._override_batch_shape),
self._zero))
self._is_maybe_batch_override = bool(
tf.get_static_value(self._override_batch_shape) is None
or tf.get_static_value(self._override_batch_shape).size != 0)
self._override_event_shape = self._maybe_validate_shape_override(
event_shape, distribution.is_scalar_event(), validate_args,
"event_shape")
self._is_event_override = prefer_static.logical_not(
prefer_static.equal(
prefer_static.rank_from_shape(self._override_event_shape),
self._zero))
self._is_maybe_event_override = bool(
tf.get_static_value(self._override_event_shape) is None
or tf.get_static_value(self._override_event_shape).size != 0)
# To convert a scalar distribution into a multivariate distribution we
# will draw dims from the sample dims, which are otherwise iid. This is
# easy to do except in the case that the base distribution has batch dims
# and we're overriding event shape. When that case happens the event dims
# will incorrectly be to the left of the batch dims. In this case we'll
# cyclically permute left the new dims.
self._needs_rotation = prefer_static.reduce_all([
self._is_event_override,
prefer_static.logical_not(self._is_batch_override),
prefer_static.logical_not(distribution.is_scalar_batch())
])
override_event_ndims = prefer_static.rank_from_shape(
self._override_event_shape)
self._rotate_ndims = _pick_scalar_condition(
self._needs_rotation, override_event_ndims, 0)
# We'll be reducing the head dims (if at all), i.e., this will be []
# if we don't need to reduce.
self._reduce_event_indices = prefer_static.range(
self._rotate_ndims - override_event_ndims, self._rotate_ndims)
self._distribution = distribution
self._bijector = bijector
super(TransformedDistribution, self).__init__(
dtype=self._distribution.dtype,
reparameterization_type=self._distribution.reparameterization_type,
validate_args=validate_args,
allow_nan_stats=self._distribution.allow_nan_stats,
parameters=parameters,
name=name)
def _inverse_log_det_jacobian(self, y):
return tf.constant(0, dtype=y.dtype)
def _forward_log_det_jacobian(self, x):
return tf.constant(0, dtype=x.dtype)
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 _event_shape_tensor(self):
return tf.constant([], dtype=tf.int32)
def _event_shape_tensor(self):
return tf.constant([self.dimension, self.dimension], dtype=tf.int32)
