def _as_distribution(self, r):
concentration = DeferredTensor(self._concentration,
lambda x: tf.cast(x, r.dtype),
dtype=r.dtype)
return tfd.Gamma(concentration=concentration,
rate=DeferredTensor(
r, lambda x: 1. / tf.math.softplus(x)))
def _as_distribution(self, r):
scale = DeferredTensor(self._scale,
lambda x: tf.cast(x, r.dtype),
dtype=r.dtype)
return tfd.LogNormal(loc=DeferredTensor(
r, lambda x: tf.math.log(tf.math.softplus(x)) - 0.5 * scale**2.),
scale=scale)
def experimental_from_mean_concentration(cls, mean, total_concentration,
**kwargs):
"""Constructs a Beta from its mean and total concentration.
**Experimental: Naming, location of this API may change.**
Total concentration, sometimes called "sample size", is the sum of the two
concentration parameters (`concentration1` and `concentration0` in
`__init__`).
Args:
mean: The mean of the constructed distribution.
total_concentration: The sum of the two concentration parameters. Must be
greater than 0.
**kwargs: Other keyword arguments passed directly to `__init__`, e.g.
`validate_args`.
Returns:
beta: A distribution with the given parameterization.
"""
dtype = dtype_util.common_dtype([mean, total_concentration],
dtype_hint=tf.float32)
total_concentration = tensor_util.convert_nonref_to_tensor(
total_concentration, dtype=dtype)
mean = tensor_util.convert_nonref_to_tensor(mean, dtype=dtype)
return cls(concentration1=DeferredTensor(
mean, lambda mean: mean * total_concentration),
concentration0=DeferredTensor(
mean, lambda mean: (1. - mean) * total_concentration),
**kwargs)
def _as_distribution(self, r):
total_count = DeferredTensor(self._total_count,
lambda x: tf.cast(x, r.dtype),
dtype=r.dtype)
return tfd.NegativeBinomial(
total_count=total_count,
logits=DeferredTensor(r, lambda x: x - tf.math.log(total_count)))
def experimental_from_mean_dispersion(cls, mean, dispersion, **kwargs):
"""Constructs a NegativeBinomial from its mean and dispersion.
**Experimental: Naming, location of this API may change.**
In this parameterization, the dispersion is defined as the reciprocal of the
total count of failures, i.e. `dispersion = 1 / total_count`.
Args:
mean: The mean of the constructed distribution.
dispersion: The reciprocal of the total_count of the constructed
distribution.
**kwargs: Other keyword arguments passed directly to `__init__`, e.g.
`validate_args`.
Returns:
neg_bin: A distribution with the given parameterization.
"""
dtype = dtype_util.common_dtype([mean, dispersion],
dtype_hint=tf.float32)
dispersion = tensor_util.convert_nonref_to_tensor(dispersion,
dtype=dtype)
mean = tensor_util.convert_nonref_to_tensor(mean, dtype=dtype)
total_count = DeferredTensor(dispersion,
tf.math.reciprocal,
dtype=dtype)
return cls(total_count=total_count,
probs=DeferredTensor(
mean, lambda mean: mean / (mean + total_count)),
**kwargs)
def _as_distribution(self, r):
total_count = DeferredTensor(self._total_count,
lambda x: tf.cast(x, r.dtype),
dtype=r.dtype)
return tfd.NegativeBinomial(
total_count=total_count,
logits=DeferredTensor(
r,
lambda x: tf.math.log( # pylint: disable=g-long-lambda
tf.math.softplus(x)) - tf.math.log(total_count)))
def experimental_from_mean_variance(cls, mean, variance, **kwargs):
"""Constructs a Beta from its mean and variance.
**Experimental: Naming, location of this API may change.**
Variance must be less than `mean * (1. - mean)`, and in particular less than
the maximal variance of 0.25, which occurs with `mean = 0.5`.
Args:
mean: The mean of the constructed distribution.
variance: The variance of the constructed distribution.
**kwargs: Other keyword arguments passed directly to `__init__`, e.g.
`validate_args`.
Returns:
beta: A distribution with the given parameterization.
"""
dtype = dtype_util.common_dtype([mean, variance],
dtype_hint=tf.float32)
variance = tensor_util.convert_nonref_to_tensor(variance, dtype=dtype)
mean = tensor_util.convert_nonref_to_tensor(mean, dtype=dtype)
total_concentration = DeferredTensor(
mean, lambda mean: mean * (1. - mean) / variance - 1.)
return cls.experimental_from_mean_concentration(
mean=mean, total_concentration=total_concentration, **kwargs)
def _default_event_space_bijector(self):
# TODO(b/145620027) Finalize choice of bijector.
deferred_scale = DeferredTensor(self.scale, lambda x: x)
return chain_bijector.Chain([
shift_bijector.Shift(
shift=deferred_scale, validate_args=self.validate_args),
softplus_bijector.Softplus(validate_args=self.validate_args)
], validate_args=self.validate_args)
def _default_event_space_bijector(self):
if tensor_util.is_ref(self.low) or tensor_util.is_ref(self.high):
scale = DeferredTensor(self.high, lambda x: x - self.low)
else:
scale = self.high - self.low
return chain_bijector.Chain([
shift_bijector.Shift(shift=self.low, validate_args=self.validate_args),
scale_bijector.Scale(scale=scale, validate_args=self.validate_args),
sigmoid_bijector.Sigmoid(validate_args=self.validate_args)
], validate_args=self.validate_args)
def experimental_from_mean_variance(cls, mean, variance, **kwargs):
"""Constructs a LogNormal from its mean and variance.
**Experimental: Naming, location of this API may change.**
Args:
mean: The mean of the constructed distribution. Must be greater than 0.
variance: The variance of the distribution. Must be greater than 0.
**kwargs: Other keyword arguments passed directly to `__init__`, e.g.
`validate_args`.
Returns:
lognormal: A distribution with the given parameterization.
"""
dtype = dtype_util.common_dtype([mean, variance], dtype_hint=tf.float32)
mean = tensor_util.convert_nonref_to_tensor(mean, dtype=dtype)
variance = tensor_util.convert_nonref_to_tensor(variance, dtype=dtype)
scale = DeferredTensor(
mean, lambda mean: tf.sqrt(tf.math.log1p(variance / mean ** 2)))
loc = DeferredTensor(
mean, lambda mean: tf.math.log(mean) - scale ** 2 / 2.)
return cls(loc=loc, scale=scale, **kwargs)
def _default_event_space_bijector(self):
# TODO(b/146568897): Resolve numerical issues by implementing a new bijector
# instead of multiplying `scale` by `(1. - 1e-6)`.
if tensor_util.is_ref(self.low) or tensor_util.is_ref(self.high):
scale = DeferredTensor(
self.high,
lambda x: (x - self.low) * (1. - 1e-6),
shape=tf.broadcast_static_shape(self.high.shape, self.low.shape))
else:
scale = (self.high - self.low) * (1. - 1e-6)
return chain_bijector.Chain([
shift_bijector.Shift(shift=self.low, validate_args=self.validate_args),
scale_bijector.Scale(scale=scale, validate_args=self.validate_args),
sigmoid_bijector.Sigmoid(validate_args=self.validate_args)
], validate_args=self.validate_args)
def experimental_from_mean_variance(cls, mean, variance, **kwargs):
"""Constructs a Gamma from its mean and variance.
**Experimental: Naming, location of this API may change.**
Args:
mean: The mean of the constructed distribution. Must be greater than 0.
variance: The variance of the distribution. Must be greater than 0.
**kwargs: Other keyword arguments passed directly to `__init__`, e.g.
`validate_args`.
Returns:
gamma: A distribution with the given parameterization.
"""
dtype = dtype_util.common_dtype([mean, variance],
dtype_hint=tf.float32)
mean = tensor_util.convert_nonref_to_tensor(mean, dtype=dtype)
variance = tensor_util.convert_nonref_to_tensor(variance, dtype=dtype)
rate = DeferredTensor(mean, lambda mean: mean / variance)
return cls(concentration=DeferredTensor(mean,
lambda mean: mean * rate),
rate=rate,
**kwargs)
def make_wishart(self, dims, new_batch_shape, old_batch_shape):
new_batch_shape_ph = (tf.constant(np.int32(new_batch_shape))
if self.is_static_shape else
tf1.placeholder_with_default(
np.int32(new_batch_shape), shape=None))
scale = self.dtype([
[[1., 0.5], [0.5, 1.]],
[[0.5, 0.25], [0.25, 0.75]],
])
scale = np.reshape(np.concatenate([scale, scale], axis=0),
old_batch_shape + [dims, dims])
scale_ph = tf1.placeholder_with_default(
scale, shape=scale.shape if self.is_static_shape else None)
wishart = tfd.WishartTriL(df=5,
scale_tril=DeferredTensor(
scale_ph, tf.linalg.cholesky),
validate_args=True)
reshape_wishart = tfd.BatchReshape(distribution=wishart,
batch_shape=new_batch_shape_ph,
validate_args=True)
return wishart, reshape_wishart
def defer_and_count_usage(var):
return DeferredTensor(usage_counting_identity, var)
def broadcasting_params(draw,
batch_shape,
params_event_ndims,
event_dim=None,
enable_vars=False,
constraint_fn_for=lambda param: identity_fn,
mutex_params=()):
"""Draws a dict of parameters which should yield the given batch shape.
The dtypes of the returned parameters are determined by their respective
constraint functions.
Args:
draw: Hypothesis MacGuffin. Supplied by `@hps.composite`.
batch_shape: A `TensorShape`. The returned parameters' batch shapes will
broadcast to this.
params_event_ndims: Python `dict` mapping the name of each parameter to a
Python `int` giving the event ndims for that parameter.
event_dim: Optional Python int giving the size of each parameter's event
dimensions (except where overridded by any applicable constraint
functions). This is shared across all parameters, permitting square event
matrices, compatible location and scale Tensors, etc. If omitted,
Hypothesis will choose one.
enable_vars: TODO(bjp): Make this `True` all the time and put variable
initialization in slicing_test. If `False`, the returned parameters are
all Tensors, never Variables or DeferredTensor.
constraint_fn_for: Python callable mapping parameter name to constraint
function. The latter is itself a Python callable which converts an
unconstrained Tensor (currently with float32 values from -200 to +200)
into one that meets the parameter's validity constraints.
mutex_params: Python iterable of Python sets. Each set gives a clique of
mutually exclusive parameters (e.g., the 'probs' and 'logits' of a
Categorical). At most one parameter from each set will appear in the
result.
Returns:
params: A Hypothesis strategy for drawing Python `dict`s mapping parameter
name to a Tensor, Variable, or DeferredTensor. The batch shapes of the
returned parameters broadcast together to the supplied `batch_shape`.
Only parameters whose names appear as keys in `params_event_ndims` will
appear (but possibly not all of them, depending on `mutex_params`).
"""
if event_dim is None:
event_dim = draw(hps.integers(min_value=2, max_value=6))
params_event_ndims = params_event_ndims or {}
remaining_params = set(params_event_ndims.keys())
params_to_use = []
while remaining_params:
param = draw(hps.one_of(map(hps.just, remaining_params)))
params_to_use.append(param)
remaining_params.remove(param)
for mutex_set in mutex_params:
if param in mutex_set:
remaining_params -= mutex_set
param_batch_shapes = draw(
broadcasting_named_shapes(batch_shape, params_to_use))
params_kwargs = dict()
for param in params_to_use:
param_batch_shape = param_batch_shapes[param]
param_event_rank = params_event_ndims[param]
param_strategy = constrained_tensors(
constraint_fn_for(param), (tensorshape_util.as_list(param_batch_shape) +
[event_dim] * param_event_rank))
params_kwargs[param] = tf.convert_to_tensor(
draw(param_strategy), dtype_hint=tf.float32, name=param)
if enable_vars and draw(hps.booleans()):
params_kwargs[param] = tf.Variable(params_kwargs[param], name=param)
alt_value = tf.convert_to_tensor(
draw(param_strategy),
dtype_hint=tf.float32,
name='{}_alt_value'.format(param))
setattr(params_kwargs[param], '_tfp_alt_value', alt_value)
if draw(hps.booleans()):
params_kwargs[param] = DeferredTensor(usage_counting_identity,
params_kwargs[param])
return params_kwargs
def _as_distribution(self, r):
return tfd.Bernoulli(logits=DeferredTensor(r, self._as_logits))
def _as_distribution(self, r):
total_count = DeferredTensor(self._total_count,
lambda x: tf.cast(x, r.dtype),
dtype=r.dtype)
return tfd.Binomial(total_count=total_count, logits=r)
def defer_and_count_usage(var): return DeferredTensor(var, make_usage_counting_identity(var))
def _as_distribution(self, r):
mean = DeferredTensor(r, self._inverse_link_fn)
return self._distribution_fn(mean)
def _defer(fn, shape, dtype):
empty = np.zeros(0)
return DeferredTensor(empty, lambda _: fn(), shape=shape, dtype=dtype)
def _as_distribution(self, r):
return tfd.Poisson(rate=DeferredTensor(r, tf.math.softplus))
def _as_distribution(self, r):
scale = DeferredTensor(self._scale,
lambda x: tf.cast(x, r.dtype),
dtype=r.dtype)
return tfd.Normal(loc=DeferredTensor(r, tf.math.reciprocal),
scale=scale)
