def semilocal_linear_trend_transition_noise(level_scale, slope_mean,
slope_scale, autoregressive_coef):
"""Build the transition noise model for a semi-local linear trend model."""
# At each timestep, the stochasticity of `level` and `slope` are given
# by `level_scale` and `slope_scale` respectively.
broadcast_batch_shape = dist_util.get_broadcast_shape(
level_scale, slope_mean, slope_scale, autoregressive_coef)
broadcast_ones = tf.ones(broadcast_batch_shape, dtype=level_scale.dtype)
scale_diag = tf.stack(
[level_scale * broadcast_ones, slope_scale * broadcast_ones], axis=-1)
# We additionally fold in a bias term implementing the nonzero `slope_mean`.
# The overall `slope` update is (from `SemiLocalLinearTrend` docstring)
# slope[t] = (slope_mean +
# autoregressive_coef * (slope[t-1] - slope_mean) +
# Normal(0., slope_scale))
# which we rewrite as
# slope[t] = (
# autoregressive_coef * slope[t-1] + # linear transition
# Normal(loc=slope_mean - autoregressive_coef * slope_mean, # noise bias
# scale=slope_scale)) # noise scale
bias = tf.stack([
tf.zeros_like(broadcast_ones), slope_mean *
(1 - autoregressive_coef) * broadcast_ones
],
axis=-1)
return tfd.MultivariateNormalDiag(loc=bias, scale_diag=scale_diag)
def _matrix(self, x1, x2):
locs = util.pad_shape_with_ones(self.locs, ndims=1, start=-2)
slopes = util.pad_shape_with_ones(self.slopes, ndims=1, start=-2)
weights_x1 = tf.math.sigmoid(
slopes *
(self.weight_fn(x1, self.feature_ndims)[..., tf.newaxis] - locs))
weights_x1 = weights_x1[..., tf.newaxis, :]
weights_x2 = tf.math.sigmoid(
slopes *
(self.weight_fn(x2, self.feature_ndims)[..., tf.newaxis] - locs))
weights_x2 = weights_x2[..., tf.newaxis, :, :]
initial_weights = (1. - weights_x1) * (1. - weights_x2)
initial_weights = tf.concat([
initial_weights,
tf.ones_like(initial_weights[..., 0])[..., tf.newaxis]
],
axis=-1)
end_weights = weights_x1 * weights_x2
end_weights = tf.concat(
[tf.ones_like(end_weights[..., 0])[..., tf.newaxis], end_weights],
axis=-1)
results = [k.matrix(x1, x2)[..., tf.newaxis] for k in self.kernels]
broadcasted_shape = distribution_util.get_broadcast_shape(*results)
results = tf.concat(
[ps.broadcast_to(r, broadcasted_shape) for r in results], axis=-1)
return tf.math.reduce_sum(initial_weights * results * end_weights,
axis=-1)
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 = tf.convert_to_tensor(loc, name="loc", dtype=dtype)
scale = tf.convert_to_tensor(scale, name="scale", dtype=dtype)
with tf.control_dependencies(
[assert_util.assert_positive(scale)] if validate_args else []):
loc = tf.identity(loc, name="loc")
scale = tf.identity(scale, name="scale")
dtype_util.assert_same_float_dtype([loc, scale])
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=loc.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 _kl_pareto_pareto(a, b, name=None):
"""Calculate the batched KL divergence KL(a || b) with a and b Pareto.
Args:
a: instance of a Pareto distribution object.
b: instance of a Pareto distribution object.
name: (optional) Name to use for created operations.
default is "kl_pareto_pareto".
Returns:
Batchwise KL(a || b)
"""
with tf.compat.v2.name_scope(name or "kl_pareto_pareto"):
# Consistent with
# http://www.mast.queensu.ca/~communications/Papers/gil-msc11.pdf, page 55
# Terminology is different from source to source for Pareto distributions.
# The 'concentration' parameter corresponds to 'a' in that source, and the
# 'scale' parameter corresponds to 'm'.
final_batch_shape = distribution_util.get_broadcast_shape(
a.concentration, b.concentration, a.scale, b.scale)
common_type = dtype_util.common_dtype(
[a.concentration, b.concentration, a.scale, b.scale], tf.float32)
return tf.where(
a.scale >= b.scale,
b.concentration * (tf.math.log(a.scale) - tf.math.log(b.scale)) +
tf.math.log(a.concentration) - tf.math.log(b.concentration) +
b.concentration / a.concentration - 1.0,
tf.broadcast_to(tf.cast(np.inf, common_type), final_batch_shape))
def _compute_flattened_covariance(self, index_points=None):
# This is of shape KN x KN, where K is the number of outputs
# Compute this explicitly via the Schur Complement of the vector kernel.
# The reason this is written explicitly as opposed to using a GPRM
# internally for reshaping is there is potential for efficiency gains when
# `observation_noise_variance = 0.`.
index_points = self._get_index_points(index_points)
kxx = self.kernel.matrix_over_all_tasks(index_points, index_points)
kxz = self.kernel.matrix_over_all_tasks(
index_points, self.observation_index_points).to_dense()
if self._observations_is_missing is not None:
kxz = tf.where(_vec(tf.math.logical_not(
self._observations_is_missing))[..., tf.newaxis, :],
kxz,
tf.zeros([], dtype=kxz.dtype))
cholinv_kzx = self.observation_cholesky.solve(kxz, adjoint_arg=True)
kxz_kzzinv_kzx = tf.linalg.matmul(
cholinv_kzx, cholinv_kzx, transpose_a=True)
flattened_covariance = kxx.to_dense() - kxz_kzzinv_kzx
if self.predictive_noise_variance is None:
return flattened_covariance
broadcast_shape = distribution_util.get_broadcast_shape(
flattened_covariance, self.predictive_noise_variance[..., tf.newaxis,
tf.newaxis])
flattened_covariance = tf.broadcast_to(flattened_covariance,
broadcast_shape)
return _add_diagonal_shift(flattened_covariance,
self.predictive_noise_variance)
def _kl_uniform_uniform(a, b, name=None):
"""Calculate the batched KL divergence KL(a || b) with a and b Uniform.
Note that the KL divergence is infinite if the support of `a` is not a subset
of the support of `b`.
Args:
a: instance of a Uniform distribution object.
b: instance of a Uniform distribution object.
name: (optional) Name to use for created operations.
default is "kl_uniform_uniform".
Returns:
Batchwise KL(a || b)
"""
with tf.name_scope(name or "kl_uniform_uniform"):
# Consistent with
# http://www.mast.queensu.ca/~communications/Papers/gil-msc11.pdf, page 60
# Watch out for the change in conventions--they use 'a' and 'b' to refer to
# lower and upper bounds respectively there.
final_batch_shape = distribution_util.get_broadcast_shape(
a.low, b.low, a.high, b.high)
dtype = dtype_util.common_dtype([a.low, a.high, b.low, b.high],
tf.float32)
return tf1.where(
(b.low <= a.low) & (a.high <= b.high),
tf.math.log(b.high - b.low) - tf.math.log(a.high - a.low),
tf.broadcast_to(
dtype_util.as_numpy_dtype(dtype)(np.inf), final_batch_shape))
def __init__(self,
concentration,
scale,
validate_args=False,
allow_nan_stats=True,
name='Weibull'):
"""Construct Weibull distributions.
The parameters `concentration` and `scale` must be shaped in a way that
supports broadcasting (e.g. `concentration + scale` is a valid operation).
Args:
concentration: Positive Float-type `Tensor`, the concentration param of the
distribution. Must contain only positive values.
scale: Positive Float-type `Tensor`, the scale param of the distribution.
Must contain only positive values.
validate_args: Python `bool` indicating whether arguments should be checked
for correctness.
allow_nan_stats: Python `bool` indicating whether nan values should be
allowed.
name: Python `str` name given to ops managed by this class.
Default value: `'Weibull'`.
Raises:
TypeError: if concentration and scale are different dtypes.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([concentration, scale],
dtype_hint=tf.float32)
concentration = tensor_util.convert_nonref_to_tensor(
concentration, name='concentration', dtype=dtype)
scale = tensor_util.convert_nonref_to_tensor(scale,
name='scale',
dtype=dtype)
# Positive scale and concentration is asserted by the incorporated
# Weibull bijector.
self._weibull_bijector = weibull_cdf_bijector.WeibullCDF(
scale=scale,
concentration=concentration,
validate_args=validate_args)
batch_shape = distribution_util.get_broadcast_shape(
concentration, scale)
super(Weibull, self).__init__(
distribution=uniform.Uniform(
# TODO(b/137665504): Use batch-adding meta-distribution to set the
# batch shape instead of tf.ones.
low=tf.zeros(batch_shape, dtype=dtype),
high=tf.ones(batch_shape, dtype=dtype),
allow_nan_stats=allow_nan_stats),
# The Weibull bijector encodes the CDF function as the forward,
# and hence needs to be inverted.
bijector=invert_bijector.Invert(self._weibull_bijector,
validate_args=validate_args),
parameters=parameters,
name=name)
def test_with_some_dynamic_shapes_works(self):
if tf.executing_eagerly(): return
x = tf.ones([2, 1, 3])
y = tf1.placeholder_with_default(np.ones([1, 5, 3], dtype=np.float32),
shape=None)
z = tf.ones([])
bcast_shape = self.evaluate(
distribution_util.get_broadcast_shape(x, y, z))
self.assertAllEqual([2, 5, 3], bcast_shape)
def test_with_some_dynamic_shapes_works(self):
x = tf.ones((2, 1, 3))
y = tf.placeholder(x.dtype)
z = tf.ones(())
with self.test_session() as sess:
bcast_shape = sess.run(
distribution_util.get_broadcast_shape(x, y, z),
feed_dict={y: np.ones((1, 5, 3)).astype(np.float32)})
self.assertAllEqual([2, 5, 3], bcast_shape)
def _compute_divisor_matrix(base_kernel, diag_shift, fixed_inputs):
"""Compute the the modified kernel with respect to the fixed inputs."""
divisor_matrix = base_kernel.matrix(fixed_inputs, fixed_inputs)
if diag_shift is not None:
diag_shift = tf.convert_to_tensor(diag_shift)
broadcast_shape = distribution_util.get_broadcast_shape(
divisor_matrix, diag_shift[..., tf.newaxis, tf.newaxis])
divisor_matrix = tf.broadcast_to(divisor_matrix, broadcast_shape)
divisor_matrix = _add_diagonal_shift(divisor_matrix,
diag_shift[..., tf.newaxis])
return divisor_matrix
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.
"""
with tf.name_scope(name, values=[loc, scale]) as name:
dtype = dtype_util.common_dtype([loc, scale], preferred_dtype=tf.float32)
loc = tf.convert_to_tensor(loc, name="loc", dtype=dtype)
scale = tf.convert_to_tensor(scale, name="scale", dtype=dtype)
with tf.control_dependencies([tf.assert_positive(scale)]
if validate_args else []):
loc = tf.identity(loc, name="loc")
scale = tf.identity(scale, name="scale")
tf.assert_same_float_dtype([loc, scale])
self._gumbel_bijector = gumbel_bijector.Gumbel(
loc=loc, scale=scale, validate_args=validate_args)
super(Gumbel, self).__init__(
distribution=uniform.Uniform(
low=tf.zeros([], dtype=loc.dtype),
high=tf.ones([], dtype=loc.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),
name=name)
def _divisor_matrix(self, fixed_inputs=None):
fixed_inputs = tf.convert_to_tensor(
self._fixed_inputs if fixed_inputs is None else fixed_inputs)
divisor_matrix = self._base_kernel.matrix(fixed_inputs, fixed_inputs)
if self._diag_shift is not None:
diag_shift = tf.convert_to_tensor(self._diag_shift)
broadcast_shape = distribution_util.get_broadcast_shape(
divisor_matrix, diag_shift[..., tf.newaxis, tf.newaxis])
divisor_matrix = tf.broadcast_to(divisor_matrix, broadcast_shape)
divisor_matrix = _add_diagonal_shift(divisor_matrix,
diag_shift[..., tf.newaxis])
return divisor_matrix
def __init__(self,
concentration1=1.,
concentration0=1.,
validate_args=False,
allow_nan_stats=True,
name='Kumaraswamy'):
"""Initialize a batch of Kumaraswamy 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)
concentration1 = tensor_util.convert_nonref_to_tensor(
concentration1, name='concentration1', dtype=dtype)
concentration0 = tensor_util.convert_nonref_to_tensor(
concentration0, name='concentration0', dtype=dtype)
self._kumaraswamy_cdf = kumaraswamy_cdf.KumaraswamyCDF(
concentration1=concentration1,
concentration0=concentration0,
validate_args=validate_args)
batch_shape = distribution_util.get_broadcast_shape(
concentration1, concentration0)
super(Kumaraswamy, self).__init__(
# TODO(b/137665504): Use batch-adding meta-distribution to set the
# batch shape instead of tf.zeros.
distribution=uniform.Uniform(
low=tf.zeros(batch_shape, dtype=dtype),
high=tf.ones([], dtype=dtype),
allow_nan_stats=allow_nan_stats),
bijector=invert.Invert(
self._kumaraswamy_cdf, validate_args=validate_args),
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.
"""
with tf.name_scope(name, values=[loc, scale]) as name:
with tf.control_dependencies(
[tf.assert_positive(scale)] if validate_args else []):
loc = tf.identity(loc, name="loc")
scale = tf.identity(scale, name="scale")
tf.assert_same_float_dtype([loc, scale])
self._gumbel_bijector = bijectors.Gumbel(
loc=loc, scale=scale, validate_args=validate_args)
super(Gumbel, self).__init__(
distribution=tf.distributions.Uniform(
low=tf.zeros([], dtype=loc.dtype),
high=tf.ones([], dtype=loc.dtype),
allow_nan_stats=allow_nan_stats),
# The Gumbel bijector encodes the quantile
# function as the forward, and hence needs to
# be inverted.
bijector=bijectors.Invert(self._gumbel_bijector),
batch_shape=distribution_util.get_broadcast_shape(loc, scale),
name=name)
def __init__(self,
concentration1=None,
concentration0=None,
validate_args=False,
allow_nan_stats=True,
name="Kumaraswamy"):
"""Initialize a batch of Kumaraswamy 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.compat.v1.name_scope(name,
values=[concentration1,
concentration0]) as name:
dtype = dtype_util.common_dtype([concentration1, concentration0],
tf.float32)
concentration1 = tf.convert_to_tensor(value=concentration1,
name="concentration1",
dtype=dtype)
concentration0 = tf.convert_to_tensor(value=concentration0,
name="concentration0",
dtype=dtype)
super(Kumaraswamy,
self).__init__(distribution=uniform.Uniform(
low=tf.zeros([], dtype=concentration1.dtype),
high=tf.ones([], dtype=concentration1.dtype),
allow_nan_stats=allow_nan_stats),
bijector=kumaraswamy_bijector.Kumaraswamy(
concentration1=concentration1,
concentration0=concentration0,
validate_args=validate_args),
batch_shape=distribution_util.get_broadcast_shape(
concentration1, concentration0),
parameters=parameters,
name=name)
def __init__(self,
concentration1=None,
concentration0=None,
validate_args=False,
allow_nan_stats=True,
name="Kumaraswamy"):
"""Initialize a batch of Kumaraswamy 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.
"""
with tf.name_scope(name, values=[concentration1, concentration0]) as name:
concentration1 = tf.convert_to_tensor(
concentration1, name="concentration1")
concentration0 = tf.convert_to_tensor(
concentration0, name="concentration0")
super(Kumaraswamy, self).__init__(
distribution=tf.distributions.Uniform(
low=tf.zeros([], dtype=concentration1.dtype),
high=tf.ones([], dtype=concentration1.dtype),
allow_nan_stats=allow_nan_stats),
bijector=bijectors.Kumaraswamy(
concentration1=concentration1,
concentration0=concentration0,
validate_args=validate_args),
batch_shape=distribution_util.get_broadcast_shape(
concentration1, concentration0),
name=name)
self._reparameterization_type = tf.distributions.FULLY_REPARAMETERIZED
def _compute_flattened_covariance(self, index_points=None):
# This is of shape KN x KN, where K is the number of outputs
index_points = self._get_index_points(index_points)
kernel_matrix = self.kernel.matrix_over_all_tasks(
index_points, index_points)
if self.observation_noise_variance is None:
return kernel_matrix
kernel_matrix = kernel_matrix.to_dense()
broadcast_shape = distribution_util.get_broadcast_shape(
kernel_matrix, self.observation_noise_variance[..., tf.newaxis,
tf.newaxis])
kernel_matrix = tf.broadcast_to(kernel_matrix, broadcast_shape)
kernel_matrix = tf.linalg.set_diag(
kernel_matrix,
tf.linalg.diag_part(kernel_matrix) +
self.observation_noise_variance[..., tf.newaxis])
kernel_matrix = tf.linalg.LinearOperatorFullMatrix(
kernel_matrix, is_non_singular=True, is_positive_definite=True)
return kernel_matrix
def _compute_covariance(self, index_points):
kernel_matrix = self.kernel.matrix(index_points, index_points)
if self._is_univariate_marginal(index_points):
# kernel_matrix thus has shape [..., 1, 1]; squeeze off the last dims and
# tack on the observation noise variance.
return (tf.squeeze(kernel_matrix, axis=[-2, -1]) +
self.observation_noise_variance)
else:
# We are compute K + obs_noise_variance * I. The shape of this matrix
# is going to be a broadcast of the shapes of K and obs_noise_variance *
# I.
broadcast_shape = distribution_util.get_broadcast_shape(
kernel_matrix,
# We pad with two single dimension since this represents a batch of
# scaled identity matrices.
self.observation_noise_variance[..., tf.newaxis, tf.newaxis])
kernel_matrix = tf.broadcast_to(kernel_matrix, broadcast_shape)
return _add_diagonal_shift(
kernel_matrix, self.observation_noise_variance[..., tf.newaxis])
def _compute_observation_scale(kernel,
observation_index_points,
cholesky_fn,
observation_noise_variance=None,
observations_is_missing=None):
"""Compute matrix square root of the kernel on observation index points."""
if observations_is_missing is not None:
observations_is_missing = tf.convert_to_tensor(observations_is_missing)
# If observations are missing, there's nothing we can do to preserve the
# operator structure, so densify.
observation_covariance = kernel.matrix_over_all_tasks(
observation_index_points, observation_index_points).to_dense()
if observation_noise_variance is not None:
broadcast_shape = distribution_util.get_broadcast_shape(
observation_covariance,
observation_noise_variance[..., tf.newaxis, tf.newaxis])
observation_covariance = tf.broadcast_to(observation_covariance,
broadcast_shape)
observation_covariance = _add_diagonal_shift(
observation_covariance, observation_noise_variance)
vec_observations_is_missing = _vec(observations_is_missing)
observation_covariance = tf.linalg.LinearOperatorFullMatrix(
psd_kernels_util.mask_matrix(
observation_covariance,
is_missing=vec_observations_is_missing),
is_non_singular=True,
is_positive_definite=True)
observation_scale = cholesky_util.cholesky_from_fn(
observation_covariance, cholesky_fn)
else:
observation_scale = mtgp._compute_flattened_scale( # pylint:disable=protected-access
kernel=kernel,
index_points=observation_index_points,
cholesky_fn=cholesky_fn,
observation_noise_variance=observation_noise_variance)
return observation_scale
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, values=[scale]) as name:
loc = tf.convert_to_tensor(loc, name="loc")
dtype = loc.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)
tf.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 partial reparameterized. loc, scale have straight
# through gradients but not the bounds.
# TODO(mfigurnov): This could be extended to use implicit gradients to
# compute derivatives for the bounds.
# https://arxiv.org/pdf/1806.01851.pdf
reparameterization_type=tf.distributions.NOT_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
graph_parents=[loc, scale, low, high],
name=name)
def _batch_shape_tensor(self):
return distribution_util.get_broadcast_shape(self.skewness,
self.tailweight, self.loc,
self.scale)
def precompute_regression_model(
kernel,
observation_index_points,
observations,
observations_is_missing=None,
index_points=None,
observation_noise_variance=None,
predictive_noise_variance=None,
mean_fn=None,
cholesky_fn=None,
validate_args=False,
allow_nan_stats=False,
name='PrecomputedMultiTaskGaussianProcessRegressionModel'):
"""Returns a MTGaussianProcessRegressionModel with precomputed quantities.
This differs from the constructor by precomputing quantities associated with
observations in a non-tape safe way. `index_points` is the only parameter
that is allowed to vary (i.e. is a `Variable` / changes after
initialization).
Specifically:
* We make `observation_index_points` and `observations` mandatory
parameters.
* We precompute `kernel(observation_index_points, observation_index_points)`
along with any other associated quantities relating to the `kernel`,
`observations` and `observation_index_points`.
A typical usecase would be optimizing kernel hyperparameters for a
`MultiTaskGaussianProcess`, and computing the posterior predictive with
respect to those optimized hyperparameters and observation / index-points
pairs.
WARNING: This method assumes `index_points` is the only varying parameter
(i.e. is a `Variable` / changes after initialization) and hence is not
tape-safe.
Args:
kernel: `PositiveSemidefiniteKernel`-like instance representing the
GP's covariance function.
observation_index_points: `float` `Tensor` representing finite collection,
or batch of collections, of points in the index set for which some data
has been observed. Shape has the form `[b1, ..., bB, e, f1, ..., fF]`
where `F` is the number of feature dimensions and must equal
`kernel.feature_ndims`, and `e` is the number (size) of index points in
each batch. `[b1, ..., bB, e]` must be broadcastable with the shape of
`observations`, and `[b1, ..., bB]` must be broadcastable with the
shapes of all other batched parameters (`kernel.batch_shape`,
`index_points`, etc). The default value is `None`, which corresponds to
the empty set of observations, and simply results in the prior
predictive model (a GP with noise of variance
`predictive_noise_variance`).
observations: `float` `Tensor` representing collection, or batch of
collections, of observations corresponding to
`observation_index_points`. Shape has the form `[b1, ..., bB, e, t]`
The batch shape `[b1, ..., bB]` must be
broadcastable with the shapes of all other batched parameters
(`kernel.batch_shape`, `index_points`, etc.). The default value is
`None`, which corresponds to the empty set of observations, and simply
results in the prior predictive model (a GP with noise of variance
`predictive_noise_variance`).
observations_is_missing: `bool` `Tensor` of shape `[..., e]`,
representing a batch of boolean masks. When `observations_is_missing`
is not `None`, the returned distribution is conditioned only on the
observations for which the corresponding elements of
`observations_is_missing` are `True`.
index_points: `float` `Tensor` representing finite collection, or batch of
collections, of points in the index set over which the GP is defined.
Shape has the form `[b1, ..., bB, e, f1, ..., fF]` where `F` is the
number of feature dimensions and must equal `kernel.feature_ndims` and
`e` is the number (size) of index points in each batch. Ultimately this
distribution corresponds to an `e`-dimensional multivariate normal. The
batch shape must be broadcastable with `kernel.batch_shape` and any
batch dims yielded by `mean_fn`.
observation_noise_variance: `float` `Tensor` representing the variance
of the noise in the Normal likelihood distribution of the model. May be
batched, in which case the batch shape must be broadcastable with the
shapes of all other batched parameters (`kernel.batch_shape`,
`index_points`, etc.).
Default value: `None`
predictive_noise_variance: `float` `Tensor` representing the variance in
the posterior predictive model. If `None`, we simply re-use
`observation_noise_variance` for the posterior predictive noise. If set
explicitly, however, we use this value. This allows us, for example, to
omit predictive noise variance (by setting this to zero) to obtain
noiseless posterior predictions of function values, conditioned on noisy
observations.
mean_fn: Python `callable` that acts on `index_points` to produce a
collection, or batch of collections, of mean values at `index_points`.
Takes a `Tensor` of shape `[b1, ..., bB, f1, ..., fF]` and returns a
`Tensor` whose shape is broadcastable with `[b1, ..., bB, t]`.
Default value: `None` implies the constant zero function.
cholesky_fn: Callable which takes a single (batch) matrix argument and
returns a Cholesky-like lower triangular factor. Default value: `None`,
in which case `make_cholesky_with_jitter_fn` is used with the `jitter`
parameter.
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: 'PrecomputedGaussianProcessRegressionModel'.
Returns
An instance of `MultiTaskGaussianProcessRegressionModel` with precomputed
quantities associated with observations.
"""
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([
index_points,
observation_index_points,
observations,
observation_noise_variance,
predictive_noise_variance,
], tf.float32)
# Convert-to-tensor arguments that are expected to not be Variables / not
# going to change.
observation_index_points = tf.convert_to_tensor(
observation_index_points, dtype=dtype)
if observation_noise_variance is not None:
observation_noise_variance = tf.convert_to_tensor(
observation_noise_variance, dtype=dtype)
observations = tf.convert_to_tensor(observations, dtype=dtype)
if observations_is_missing is not None:
observations_is_missing = tf.convert_to_tensor(
observations_is_missing)
if cholesky_fn is None:
cholesky_fn = cholesky_util.make_cholesky_with_jitter_fn()
else:
if not callable(cholesky_fn):
raise ValueError('`cholesky_fn` must be a Python callable')
if mean_fn is None:
mean_fn = lambda x: tf.zeros([1], dtype=dtype)
else:
if not callable(mean_fn):
raise ValueError('`mean_fn` must be a Python callable')
if observations_is_missing is not None:
# If observations are missing, there's nothing we can do to preserve the
# operator structure, so densify.
observation_covariance = kernel.matrix_over_all_tasks(
observation_index_points,
observation_index_points).to_dense()
if observation_noise_variance is not None:
broadcast_shape = distribution_util.get_broadcast_shape(
observation_covariance,
observation_noise_variance[..., tf.newaxis,
tf.newaxis])
observation_covariance = tf.broadcast_to(
observation_covariance, broadcast_shape)
observation_covariance = _add_diagonal_shift(
observation_covariance, observation_noise_variance)
vec_observations_is_missing = _vec(observations_is_missing)
observation_covariance = tf.linalg.LinearOperatorFullMatrix(
psd_kernels_util.mask_matrix(
observation_covariance,
is_missing=vec_observations_is_missing),
is_non_singular=True,
is_positive_definite=True)
observation_scale = cholesky_util.cholesky_from_fn(
observation_covariance, cholesky_fn)
else:
observation_scale = mtgp._compute_flattened_scale( # pylint:disable=protected-access
kernel=kernel,
index_points=observation_index_points,
cholesky_fn=cholesky_fn,
observation_noise_variance=observation_noise_variance)
# Note that the conditional mean is
# k(x, o) @ (k(o, o) + sigma**2)^-1 obs. We can precompute the latter
# term since it won't change per iteration.
vec_diff = _vec(observations - mean_fn(observation_index_points))
if observations_is_missing is not None:
vec_diff = tf.where(vec_observations_is_missing,
tf.zeros([], dtype=vec_diff.dtype),
vec_diff)
solve_on_observations = observation_scale.solvevec(
observation_scale.solvevec(vec_diff), adjoint=True)
def flattened_conditional_mean_fn(x):
return _flattened_conditional_mean_fn_helper(
x,
kernel,
observations,
observation_index_points,
observations_is_missing,
observation_scale,
mean_fn,
solve_on_observations=solve_on_observations)
mtgprm = MultiTaskGaussianProcessRegressionModel(
kernel=kernel,
observation_index_points=observation_index_points,
observations=observations,
index_points=index_points,
observation_noise_variance=observation_noise_variance,
predictive_noise_variance=predictive_noise_variance,
cholesky_fn=cholesky_fn,
_flattened_conditional_mean_fn=flattened_conditional_mean_fn,
_observation_scale=observation_scale,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
name=name)
return mtgprm
def __init__(self,
loc,
scale,
validate_args=False,
allow_nan_stats=True,
name='Moyal'):
"""Construct Moyal 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: `'Moyal'`.
Raises:
TypeError: if loc and scale are different dtypes.
#### References
[1] J.E. Moyal, "XXX. Theory of ionization fluctuations",
The London, Edinburgh, and Dublin Philosophical Magazine
and Journal of Science.
https://www.tandfonline.com/doi/abs/10.1080/14786440308521076
[2] G. Cordeiro, J. Nobre, R. Pescim, E. Ortega,
"The beta Moyal: a useful skew distribution",
https://www.arpapress.com/Volumes/Vol10Issue2/IJRRAS_10_2_02.pdf
"""
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 Moyal bijector.
self._moyal_bijector = moyal_cdf_bijector.MoyalCDF(
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.
batch_shape = distribution_util.get_broadcast_shape(loc, scale)
super(Moyal, self).__init__(
# TODO(b/137665504): Use batch-adding meta-distribution to set the
# batch shape instead of tf.ones.
distribution=uniform.Uniform(low=np.finfo(
dtype_util.as_numpy_dtype(dtype)).tiny,
high=tf.ones(batch_shape,
dtype=dtype),
allow_nan_stats=allow_nan_stats),
# The Moyal bijector encodes the CDF function as the forward,
# and hence needs to be inverted.
bijector=invert_bijector.Invert(self._moyal_bijector,
validate_args=validate_args),
parameters=parameters,
name=name)
def __init__(self,
loc,
scale,
skewness=None,
tailweight=None,
distribution=None,
validate_args=False,
allow_nan_stats=True,
name="SinhArcsinh"):
"""Construct SinhArcsinh distribution on `(-inf, inf)`.
Arguments `(loc, scale, skewness, tailweight)` must have broadcastable shape
(indexing batch dimensions). They must all have the same `dtype`.
Args:
loc: Floating-point `Tensor`.
scale: `Tensor` of same `dtype` as `loc`.
skewness: Skewness parameter. Default is `0.0` (no skew).
tailweight: Tailweight parameter. Default is `1.0` (unchanged tailweight)
distribution: `tf.Distribution`-like instance. Distribution that is
transformed to produce this distribution.
Default is `tfd.Normal(0., 1.)`.
Must be a scalar-batch, scalar-event distribution. Typically
`distribution.reparameterization_type = FULLY_REPARAMETERIZED` or it is
a function of non-trainable parameters. WARNING: If you backprop through
a `SinhArcsinh` sample and `distribution` is not
`FULLY_REPARAMETERIZED` yet is a function of trainable variables, then
the gradient will be incorrect!
validate_args: Python `bool`, default `False`. When `True` distribution
parameters are checked for validity despite possibly degrading runtime
performance. When `False` invalid inputs may silently render incorrect
outputs.
allow_nan_stats: Python `bool`, default `True`. When `True`,
statistics (e.g., mean, mode, variance) use the value "`NaN`" to
indicate the result is undefined. When `False`, an exception is raised
if one or more of the statistic's batch members are undefined.
name: Python `str` name prefixed to Ops created by this class.
"""
parameters = dict(locals())
with tf.compat.v2.name_scope(name) as name:
dtype = dtype_util.common_dtype([loc, scale, skewness, tailweight],
tf.float32)
loc = tf.convert_to_tensor(value=loc, name="loc", dtype=dtype)
scale = tf.convert_to_tensor(value=scale,
name="scale",
dtype=dtype)
tailweight = 1. if tailweight is None else tailweight
has_default_skewness = skewness is None
skewness = 0. if skewness is None else skewness
tailweight = tf.convert_to_tensor(value=tailweight,
name="tailweight",
dtype=dtype)
skewness = tf.convert_to_tensor(value=skewness,
name="skewness",
dtype=dtype)
batch_shape = distribution_util.get_broadcast_shape(
loc, scale, tailweight, 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)
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:
loc = distribution_util.with_dependencies(asserts, loc)
# Make the SAS bijector, 'F'.
f = sinh_arcsinh_bijector.SinhArcsinh(skewness=skewness,
tailweight=tailweight)
if has_default_skewness:
f_noskew = f
else:
f_noskew = sinh_arcsinh_bijector.SinhArcsinh(
skewness=skewness.dtype.as_numpy_dtype(0.),
tailweight=tailweight)
# Make the AffineScalar bijector, Z --> loc + scale * Z (2 / F_0(2))
c = 2 * scale / f_noskew.forward(
tf.convert_to_tensor(value=2, dtype=dtype))
affine = affine_scalar_bijector.AffineScalar(
shift=loc, scale=c, 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
self._loc = loc
self._scale = scale
self._tailweight = tailweight
self._skewness = skewness
def __init__(self,
num_timesteps,
level_scale,
slope_scale,
initial_state_prior,
observation_noise_scale=0.,
initial_step=0,
validate_args=False,
allow_nan_stats=True,
name=None):
"""Build a state space model implementing a local linear trend.
Args:
num_timesteps: Scalar `int` `Tensor` number of timesteps to model
with this distribution.
level_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
level transitions.
slope_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
slope transitions.
initial_state_prior: instance of `tfd.MultivariateNormal`
representing the prior distribution on latent states; must
have event shape `[2]`.
observation_noise_scale: Scalar (any additional dimensions are
treated as batch dimensions) `float` `Tensor` indicating the standard
deviation of the observation noise.
initial_step: Optional scalar `int` `Tensor` specifying the starting
timestep.
Default value: 0.
validate_args: Python `bool`. Whether to validate input
with asserts. If `validate_args` is `False`, and the inputs are
invalid, correct behavior is not guaranteed.
Default value: `False`.
allow_nan_stats: Python `bool`. 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.
Default value: `True`.
name: Python `str` name prefixed to ops created by this class.
Default value: "LocalLinearTrendStateSpaceModel".
"""
with tf.name_scope(name, 'LocalLinearTrendStateSpaceModel',
[level_scale, slope_scale]) as name:
# The initial state prior determines the dtype of sampled values.
# Other model parameters must have the same dtype.
dtype = initial_state_prior.dtype
level_scale = tf.convert_to_tensor(
level_scale, name='level_scale', dtype=dtype)
slope_scale = tf.convert_to_tensor(
slope_scale, name='slope_scale', dtype=dtype)
observation_noise_scale = tf.convert_to_tensor(
observation_noise_scale, name='observation_noise_scale', dtype=dtype)
# Explicitly broadcast all parameters to the same batch shape. This
# allows us to use `tf.stack` for a compact model specification.
broadcast_batch_shape = dist_util.get_broadcast_shape(
level_scale, slope_scale)
broadcast_ones = tf.ones(broadcast_batch_shape, dtype=dtype)
self._level_scale = level_scale
self._slope_scale = slope_scale
self._observation_noise_scale = observation_noise_scale
# Construct a linear Gaussian state space model implementing the
# local linear trend model. See "Mathematical Details" in the
# class docstring for further explanation.
super(LocalLinearTrendStateSpaceModel, self).__init__(
num_timesteps=num_timesteps,
transition_matrix=tf.constant(
[[1., 1.], [0., 1.]], dtype=dtype, name='transition_matrix'),
transition_noise=tfd.MultivariateNormalDiag(
scale_diag=tf.stack(
[level_scale * broadcast_ones, slope_scale * broadcast_ones],
axis=-1),
name='transition_noise'),
observation_matrix=tf.constant(
[[1., 0.]], dtype=dtype, name='observation_matrix'),
observation_noise=tfd.MultivariateNormalDiag(
scale_diag=observation_noise_scale[..., tf.newaxis],
name='observation_noise'),
initial_state_prior=initial_state_prior,
initial_step=initial_step,
allow_nan_stats=allow_nan_stats,
validate_args=validate_args,
name=name)
def _compute_flattened_scale(kernel,
index_points,
cholesky_fn,
observation_noise_variance=None):
"""Computes a matrix square root of the flattened covariance matrix.
Given a multi-task kernel `k`, computes a matrix square root of the
matrix over all tasks of `index_points`. That is, compute `S` such that
`S^T @ S = k.matrix_over_all_tasks(index_points, index_points)`.
In the case of a `Separable` or `Independent` kernel, this function tries to
do this efficiently in O(N^3 + T^3) time where `N` is the number of
`index_points` and `T` is the number of tasks.
Args:
kernel: `MultiTaskKernel`-like instance representing the GP's covariance
function.
index_points: `float` `Tensor` representing finite collection, or batch of
collections, of points in the index set over which the GP is defined.
Shape has the form `[b1, ..., bB, e, f1, ..., fF]` where `F` is the
number of feature dimensions and must equal `kernel.feature_ndims` and
`e` is the number (size) of index points in each batch. Ultimately this
distribution corresponds to an `e`-dimensional multivariate normal. The
batch shape must be broadcastable with `kernel.batch_shape`.
cholesky_fn: Callable which takes a single (batch) matrix argument and
returns a Cholesky-like lower triangular factor. Default value: `None`,
in which case `make_cholesky_with_jitter_fn(1e-6)` is used.
observation_noise_variance: `float` `Tensor` representing the variance
of the noise in the Normal likelihood distribution of the model. May be
batched, in which case the batch shape must be broadcastable with the
shapes of all other batched parameters (`kernel.batch_shape`,
`index_points`, etc.).
Default value: `None`
Returns:
scale_operator: `LinearOperator` representing a matrix square root of
the flattened kernel matrix over all tasks.
"""
# This is of shape KN x KN, where K is the number of outputs
kernel_matrix = kernel.matrix_over_all_tasks(index_points, index_points)
if observation_noise_variance is None:
return cholesky_util.cholesky_from_fn(kernel_matrix, cholesky_fn)
observation_noise_variance = tf.convert_to_tensor(
observation_noise_variance)
# We can add the observation noise to each block.
if isinstance(kernel, multitask_kernel.Independent):
# The Independent kernel matrix is realized as a kronecker product of the
# kernel over inputs, and an identity matrix per task (representing
# independent tasks). Update the diagonal of the first matrix and take the
# cholesky of it (since the cholesky of the second matrix will remain the
# identity matrix.)
base_kernel_matrix = kernel_matrix.operators[0].to_dense()
broadcast_shape = distribution_util.get_broadcast_shape(
base_kernel_matrix, observation_noise_variance[..., tf.newaxis,
tf.newaxis])
base_kernel_matrix = tf.broadcast_to(base_kernel_matrix,
broadcast_shape)
base_kernel_matrix = tf.linalg.set_diag(
base_kernel_matrix,
tf.linalg.diag_part(base_kernel_matrix) +
observation_noise_variance[..., tf.newaxis])
base_kernel_matrix = tf.linalg.LinearOperatorFullMatrix(
base_kernel_matrix)
kernel_matrix = tf.linalg.LinearOperatorKronecker(
operators=[base_kernel_matrix] + kernel_matrix.operators[1:])
return cholesky_util.cholesky_from_fn(kernel_matrix, cholesky_fn)
if isinstance(kernel, multitask_kernel.Separable):
# When `kernel_matrix` is a kronecker product, we can compute
# an eigenvalue decomposition to get a matrix square-root, which will
# be faster than densifying the kronecker product.
# Let K = A X B. Let A (and B) have an eigenvalue decomposition of
# U @ D @ U^T, where U is an orthogonal matrix. Then,
# K = (U_A @ D_A @ U_A^T) X (U_B @ D_B @ U_B^T) =
# (U_A X U_B) @ (D_A X D_B) @ (U_A X U_B)^T
# Thus, a matrix square root of K would be
# (U_A X U_B) @ (sqrt(D_A) X sqrt(D_B)) which offers
# efficient matmul and solves.
# Now, if we update the diagonal by `v * I`, we have
# (U_A X U_B) @ (sqrt((D_A X D_B + vI)) @ (U_A X U_B)^T
# which still admits an efficient matmul and solve.
kronecker_diags = []
kronecker_orths = []
for block in kernel_matrix.operators:
diag, orth = tf.linalg.eigh(block.to_dense())
kronecker_diags.append(tf.linalg.LinearOperatorDiag(diag))
kronecker_orths.append(
linear_operator_unitary.LinearOperatorUnitary(orth))
full_diag = tf.linalg.LinearOperatorKronecker(
kronecker_diags).diag_part()
full_diag = full_diag + observation_noise_variance[..., tf.newaxis]
scale_diag = tf.math.sqrt(full_diag)
diag_operator = tf.linalg.LinearOperatorDiag(scale_diag,
is_square=True,
is_non_singular=True,
is_positive_definite=True)
orthogonal_operator = tf.linalg.LinearOperatorKronecker(
kronecker_orths, is_square=True, is_non_singular=True)
# This is efficient as a scale matrix. When used for matmuls, we take
# advantage of the kronecker product and diagonal operator. When used for
# solves, we take advantage of the orthogonal and diagonal structure,
# which essentially reduces to the matmul case.
return orthogonal_operator.matmul(diag_operator)
# By default densify the kernel matrix and add noise.
kernel_matrix = kernel_matrix.to_dense()
broadcast_shape = distribution_util.get_broadcast_shape(
kernel_matrix, observation_noise_variance[..., tf.newaxis, tf.newaxis])
kernel_matrix = tf.broadcast_to(kernel_matrix, broadcast_shape)
kernel_matrix = tf.linalg.set_diag(
kernel_matrix,
tf.linalg.diag_part(kernel_matrix) +
observation_noise_variance[..., tf.newaxis])
kernel_matrix = tf.linalg.LinearOperatorFullMatrix(kernel_matrix)
kernel_cholesky = cholesky_util.cholesky_from_fn(kernel_matrix,
cholesky_fn)
return kernel_cholesky
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 `tf.distributions.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, values=[loc, scale, skewness, tailweight]) as name:
loc = tf.convert_to_tensor(loc, name="loc")
dtype = loc.dtype
scale = tf.convert_to_tensor(scale, name="scale", dtype=dtype)
tailweight = 1. if tailweight is None else tailweight
has_default_skewness = skewness is None
skewness = 0. if skewness is None else skewness
tailweight = tf.convert_to_tensor(
tailweight, name="tailweight", dtype=dtype)
skewness = tf.convert_to_tensor(skewness, name="skewness", dtype=dtype)
batch_shape = distribution_util.get_broadcast_shape(
loc, scale, tailweight, 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)
if distribution is None:
distribution = tf.distributions.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:
loc = control_flow_ops.with_dependencies(asserts, loc)
# Make the SAS bijector, 'F'.
f = bijectors.SinhArcsinh(
skewness=skewness, tailweight=tailweight)
if has_default_skewness:
f_noskew = f
else:
f_noskew = bijectors.SinhArcsinh(
skewness=skewness.dtype.as_numpy_dtype(0.),
tailweight=tailweight)
# Make the AffineScalar bijector, Z --> loc + scale * Z (2 / F_0(2))
c = 2 * scale / f_noskew.forward(tf.convert_to_tensor(2, dtype=dtype))
affine = bijectors.AffineScalar(
shift=loc,
scale=c,
validate_args=validate_args)
bijector = bijectors.Chain([affine, f])
super(SinhArcsinh, self).__init__(
distribution=distribution,
bijector=bijector,
batch_shape=batch_shape,
validate_args=validate_args,
name=name)
self._parameters = parameters
self._loc = loc
self._scale = scale
self._tailweight = tailweight
self._skewness = skewness
def __init__(self,
kernel,
observation_index_points,
observations,
observations_is_missing=None,
index_points=None,
mean_fn=None,
observation_noise_variance=None,
predictive_noise_variance=None,
cholesky_fn=None,
validate_args=False,
allow_nan_stats=False,
name='MultiTaskGaussianProcessRegressionModelWithCholesky'):
"""Construct a MultiTaskGaussianProcessRegressionModelWithCholesky instance.
WARNING: This method assumes `index_points` is the only varying parameter
(i.e. is a `Variable` / changes after initialization) and hence is not
tape-safe.
Args:
kernel: `MultiTaskKernel`-like instance representing the GP's covariance
function.
observation_index_points: `float` `Tensor` representing finite collection,
or batch of collections, of points in the index set for which some data
has been observed. Shape has the form `[b1, ..., bB, e, f1, ..., fF]`
where `F` is the number of feature dimensions and must equal
`kernel.feature_ndims`, and `e` is the number (size) of index points in
each batch. `[b1, ..., bB, e]` must be broadcastable with the shape of
`observations`, and `[b1, ..., bB]` must be broadcastable with the
shapes of all other batched parameters (`kernel.batch_shape`,
`index_points`, etc).
observations: `float` `Tensor` representing collection, or batch of
collections, of observations corresponding to
`observation_index_points`. Shape has the form `[b1, ..., bB, e, t]`,
which must be broadcastable with the batch and example shapes of
`observation_index_points`. The batch shape `[b1, ..., bB]` must be
broadcastable with the shapes of all other batched parameters
(`kernel.batch_shape`, `index_points`, etc.).
observations_is_missing: `bool` `Tensor` of shape `[..., e, t]`,
representing a batch of boolean masks. When
`observations_is_missing` is not `None`, this distribution is
conditioned only on the observations for which the
corresponding elements of `observations_is_missing` are `False`.
index_points: `float` `Tensor` representing finite collection, or batch of
collections, of points in the index set over which the GP is defined.
Shape has the form `[b1, ..., bB, e, f1, ..., fF]` where `F` is the
number of feature dimensions and must equal `kernel.feature_ndims` and
`e` is the number (size) of index points in each batch. Ultimately this
distribution corresponds to an `e`-dimensional multivariate normal. The
batch shape must be broadcastable with `kernel.batch_shape`.
mean_fn: Python `callable` that acts on `index_points` to produce a (batch
of) collection of mean values at `index_points`. Takes a `Tensor` of
shape `[b1, ..., bB, e, f1, ..., fF]` and returns a `Tensor` whose shape
is broadcastable with `[b1, ..., bB, e, t]`, where `t` is the number of
tasks.
observation_noise_variance: `float` `Tensor` representing the variance of
the noise in the Normal likelihood distribution of the model. May be
batched, in which case the batch shape must be broadcastable with the
shapes of all other batched parameters (`kernel.batch_shape`,
`index_points`, etc.).
Default value: `None`
predictive_noise_variance: `float` `Tensor` representing the variance in
the posterior predictive model. If `None`, we simply re-use
`observation_noise_variance` for the posterior predictive noise. If set
explicitly, however, we use this value. This allows us, for example, to
omit predictive noise variance (by setting this to zero) to obtain
noiseless posterior predictions of function values, conditioned on noisy
observations.
cholesky_fn: Callable which takes a single (batch) matrix argument and
returns a Cholesky-like lower triangular factor. Default value: `None`,
in which case `make_cholesky_with_jitter_fn(1e-6)` is used.
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: 'MultiTaskGaussianProcessRegressionModel'.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
if not isinstance(kernel, multitask_kernel.MultiTaskKernel):
raise ValueError('`kernel` must be a `MultiTaskKernel`.')
dtype = dtype_util.common_dtype([
index_points, observation_index_points, observations,
observation_noise_variance, predictive_noise_variance
], tf.float32)
index_points = tensor_util.convert_nonref_to_tensor(
index_points, dtype=dtype, name='index_points')
observation_index_points = tf.convert_to_tensor(
observation_index_points,
dtype=dtype,
name='observation_index_points')
observations = tf.convert_to_tensor(
observations, dtype=dtype, name='observations')
if observations_is_missing is not None:
observations_is_missing = tf.convert_to_tensor(
observations_is_missing, dtype=tf.bool)
if observation_noise_variance is not None:
observation_noise_variance = tf.convert_to_tensor(
observation_noise_variance,
dtype=dtype,
name='observation_noise_variance')
predictive_noise_variance = tensor_util.convert_nonref_to_tensor(
predictive_noise_variance,
dtype=dtype,
name='predictive_noise_variance')
if predictive_noise_variance is None:
predictive_noise_variance = observation_noise_variance
if cholesky_fn is None:
self._cholesky_fn = cholesky_util.make_cholesky_with_jitter_fn()
else:
if not callable(cholesky_fn):
raise ValueError('`cholesky_fn` must be a Python callable')
self._cholesky_fn = cholesky_fn
self._kernel = kernel
self._index_points = index_points
# Scalar or vector the size of the number of tasks.
if mean_fn is not None:
if not callable(mean_fn):
raise ValueError('`mean_fn` must be a Python callable')
self._mean_fn = mean_fn
self._observation_noise_variance = observation_noise_variance
self._predictive_noise_variance = predictive_noise_variance
self._index_ponts = index_points
self._observation_index_points = observation_index_points
self._observations = observations
self._observations_is_missing = observations_is_missing
observation_covariance = self.kernel.matrix_over_all_tasks(
observation_index_points, observation_index_points)
if observation_noise_variance is not None:
observation_covariance = observation_covariance.to_dense()
broadcast_shape = distribution_util.get_broadcast_shape(
observation_covariance, observation_noise_variance[..., tf.newaxis,
tf.newaxis])
observation_covariance = tf.broadcast_to(observation_covariance,
broadcast_shape)
observation_covariance = _add_diagonal_shift(observation_covariance,
observation_noise_variance)
observation_covariance = tf.linalg.LinearOperatorFullMatrix(
observation_covariance,
is_non_singular=True,
is_positive_definite=True)
if observations_is_missing is not None:
vec_observations_is_missing = _vec(observations_is_missing)
observation_covariance = tf.linalg.LinearOperatorFullMatrix(
psd_kernels_util.mask_matrix(
observation_covariance.to_dense(),
mask=~vec_observations_is_missing),
is_non_singular=True,
is_positive_definite=True)
self._observation_cholesky = cholesky_util.cholesky_from_fn(
observation_covariance, self._cholesky_fn)
# Note that the conditional mean is
# k(x, o) @ (k(o, o) + sigma**2)^-1 obs. We can precompute the latter
# term since it won't change per iteration.
if mean_fn:
vec_observations = _vec(observations -
mean_fn(observation_index_points))
else:
vec_observations = _vec(observations)
if observations_is_missing is not None:
vec_observations = tf.where(~vec_observations_is_missing,
vec_observations,
tf.zeros([], dtype=vec_observations.dtype))
self._solve_on_obs = self._observation_cholesky.solvevec(
self._observation_cholesky.solvevec(vec_observations), adjoint=True)
super(MultiTaskGaussianProcessRegressionModel, 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,
num_timesteps,
level_scale,
slope_scale,
initial_state_prior,
observation_noise_scale=0.,
initial_step=0,
validate_args=False,
allow_nan_stats=True,
name=None):
"""Build a state space model implementing a local linear trend.
Args:
num_timesteps: Scalar `int` `Tensor` number of timesteps to model
with this distribution.
level_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
level transitions.
slope_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
slope transitions.
initial_state_prior: instance of `tfd.MultivariateNormal`
representing the prior distribution on latent states; must
have event shape `[2]`.
observation_noise_scale: Scalar (any additional dimensions are
treated as batch dimensions) `float` `Tensor` indicating the standard
deviation of the observation noise.
initial_step: Optional scalar `int` `Tensor` specifying the starting
timestep.
Default value: 0.
validate_args: Python `bool`. Whether to validate input
with asserts. If `validate_args` is `False`, and the inputs are
invalid, correct behavior is not guaranteed.
Default value: `False`.
allow_nan_stats: Python `bool`. 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.
Default value: `True`.
name: Python `str` name prefixed to ops created by this class.
Default value: "LocalLinearTrendStateSpaceModel".
"""
with tf.name_scope(name, 'LocalLinearTrendStateSpaceModel',
[level_scale, slope_scale]) as name:
# The initial state prior determines the dtype of sampled values.
# Other model parameters must have the same dtype.
dtype = initial_state_prior.dtype
level_scale = tf.convert_to_tensor(
level_scale, name='level_scale', dtype=dtype)
slope_scale = tf.convert_to_tensor(
slope_scale, name='slope_scale', dtype=dtype)
observation_noise_scale = tf.convert_to_tensor(
observation_noise_scale, name='observation_noise_scale', dtype=dtype)
# Explicitly broadcast all parameters to the same batch shape. This
# allows us to use `tf.stack` for a compact model specification.
broadcast_batch_shape = dist_util.get_broadcast_shape(
level_scale, slope_scale)
broadcast_ones = tf.ones(broadcast_batch_shape, dtype=dtype)
self._level_scale = level_scale
self._slope_scale = slope_scale
self._observation_noise_scale = observation_noise_scale
# Construct a linear Gaussian state space model implementing the
# local linear trend model. See "Mathematical Details" in the
# class docstring for further explanation.
super(LocalLinearTrendStateSpaceModel, self).__init__(
num_timesteps=num_timesteps,
transition_matrix=tf.constant(
[[1., 1.], [0., 1.]], dtype=dtype, name='transition_matrix'),
transition_noise=tfd.MultivariateNormalDiag(
scale_diag=tf.stack(
[level_scale * broadcast_ones, slope_scale * broadcast_ones],
axis=-1),
name='transition_noise'),
observation_matrix=tf.constant(
[[1., 0.]], dtype=dtype, name='observation_matrix'),
observation_noise=tfd.MultivariateNormalDiag(
scale_diag=observation_noise_scale[..., tf.newaxis],
name='observation_noise'),
initial_state_prior=initial_state_prior,
initial_step=initial_step,
allow_nan_stats=allow_nan_stats,
validate_args=validate_args,
name=name)
def __init__(self,
num_timesteps,
level_scale,
slope_scale,
initial_state_prior,
observation_noise_scale=0.,
name=None,
**linear_gaussian_ssm_kwargs):
"""Build a state space model implementing a local linear trend.
Args:
num_timesteps: Scalar `int` `Tensor` number of timesteps to model
with this distribution.
level_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
level transitions.
slope_scale: Scalar (any additional dimensions are treated as batch
dimensions) `float` `Tensor` indicating the standard deviation of the
slope transitions.
initial_state_prior: instance of `tfd.MultivariateNormal`
representing the prior distribution on latent states; must
have event shape `[2]`.
observation_noise_scale: Scalar (any additional dimensions are
treated as batch dimensions) `float` `Tensor` indicating the standard
deviation of the observation noise.
name: Python `str` name prefixed to ops created by this class.
Default value: "LocalLinearTrendStateSpaceModel".
**linear_gaussian_ssm_kwargs: Optional additional keyword arguments to
to the base `tfd.LinearGaussianStateSpaceModel` constructor.
"""
parameters = dict(locals())
parameters.update(linear_gaussian_ssm_kwargs)
del parameters['linear_gaussian_ssm_kwargs']
with tf.name_scope(name or 'LocalLinearTrendStateSpaceModel') as name:
# The initial state prior determines the dtype of sampled values.
# Other model parameters must have the same dtype.
dtype = initial_state_prior.dtype
level_scale = tf.convert_to_tensor(value=level_scale,
name='level_scale',
dtype=dtype)
slope_scale = tf.convert_to_tensor(value=slope_scale,
name='slope_scale',
dtype=dtype)
observation_noise_scale = tf.convert_to_tensor(
value=observation_noise_scale,
name='observation_noise_scale',
dtype=dtype)
# Explicitly broadcast all parameters to the same batch shape. This
# allows us to use `tf.stack` for a compact model specification.
broadcast_batch_shape = dist_util.get_broadcast_shape(
level_scale, slope_scale)
broadcast_ones = tf.ones(broadcast_batch_shape, dtype=dtype)
self._level_scale = level_scale
self._slope_scale = slope_scale
self._observation_noise_scale = observation_noise_scale
# Construct a linear Gaussian state space model implementing the
# local linear trend model. See "Mathematical Details" in the
# class docstring for further explanation.
super(LocalLinearTrendStateSpaceModel, self).__init__(
num_timesteps=num_timesteps,
transition_matrix=tf.constant([[1., 1.], [0., 1.]],
dtype=dtype,
name='transition_matrix'),
transition_noise=tfd.MultivariateNormalDiag(
scale_diag=tf.stack([
level_scale * broadcast_ones,
slope_scale * broadcast_ones
],
axis=-1),
name='transition_noise'),
observation_matrix=tf.constant([[1., 0.]],
dtype=dtype,
name='observation_matrix'),
observation_noise=tfd.MultivariateNormalDiag(
scale_diag=observation_noise_scale[..., tf.newaxis],
name='observation_noise'),
initial_state_prior=initial_state_prior,
name=name,
**linear_gaussian_ssm_kwargs)
self._parameters = parameters
def __init__(self,
base_kernel,
fixed_inputs,
diag_shift=None,
validate_args=False,
name='SchurComplement'):
"""Construct a SchurComplement kernel instance.
Args:
base_kernel: A `PositiveSemidefiniteKernel` instance, the kernel used to
build the block matrices of which this kernel computes the Schur
complement.
fixed_inputs: A Tensor, representing a collection of inputs. The Schur
complement that this kernel computes comes from a block matrix, whose
bottom-right corner is derived from `base_kernel.matrix(fixed_inputs,
fixed_inputs)`, and whose top-right and bottom-left pieces are
constructed by computing the base_kernel at pairs of input locations
together with these `fixed_inputs`. `fixed_inputs` is allowed to be an
empty collection (either `None` or having a zero shape entry), in which
case the kernel falls back to the trivial application of `base_kernel`
to inputs. See class-level docstring for more details on the exact
computation this does; `fixed_inputs` correspond to the `Z` structure
discussed there. `fixed_inputs` is assumed to have shape `[b1, ..., bB,
N, f1, ..., fF]` where the `b`'s are batch shape entries, the `f`'s are
feature_shape entries, and `N` is the number of fixed inputs. Use of
this kernel entails a 1-time O(N^3) cost of computing the Cholesky
decomposition of the k(Z, Z) matrix. The batch shape elements of
`fixed_inputs` must be broadcast compatible with
`base_kernel.batch_shape`.
diag_shift: A floating point scalar to be added to the diagonal of the
divisor_matrix before computing its Cholesky.
validate_args: If `True`, parameters are checked for validity despite
possibly degrading runtime performance.
Default value: `False`
name: Python `str` name prefixed to Ops created by this class.
Default value: `"SchurComplement"`
"""
with tf.compat.v1.name_scope(name, values=[base_kernel,
fixed_inputs]) as name:
# If the base_kernel doesn't have a specified dtype, we can't pass it off
# to common_dtype, which always expects `tf.as_dtype(dtype)` to work (and
# it doesn't if the given `dtype` is None.
# TODO(b/130421035): Consider changing common_dtype to allow Nones, and
# clean this up after.
#
# Thus, we spell out the logic
# here: use the dtype of `fixed_inputs` if possible. If base_kernel.dtype
# is not None, use the usual logic.
if base_kernel.dtype is None:
dtype = None if fixed_inputs is None else fixed_inputs.dtype
else:
dtype = dtype_util.common_dtype([base_kernel, fixed_inputs],
tf.float32)
self._base_kernel = base_kernel
self._fixed_inputs = (None if fixed_inputs is None else
tf.convert_to_tensor(value=fixed_inputs,
dtype=dtype))
if not self._is_fixed_inputs_empty():
# We create and store this matrix here, so that we get the caching
# benefit when we later access its cholesky. If we computed the matrix
# every time we needed the cholesky, the bijector cache wouldn't be hit.
self._divisor_matrix = base_kernel.matrix(
fixed_inputs, fixed_inputs)
if diag_shift is not None:
broadcast_shape = distribution_util.get_broadcast_shape(
self._divisor_matrix, diag_shift[..., tf.newaxis])
self._divisor_matrix = tf.broadcast_to(
self._divisor_matrix, broadcast_shape)
self._divisor_matrix = _add_diagonal_shift(
self._divisor_matrix, diag_shift)
self._cholesky_bijector = invert.Invert(
cholesky_outer_product.CholeskyOuterProduct())
super(SchurComplement, self).__init__(base_kernel.feature_ndims,
dtype=dtype,
name=name)
def __init__(self,
loc,
scale,
concentration,
validate_args=False,
allow_nan_stats=True,
name='GeneralizedExtremeValue'):
"""Construct generalized extreme value distribution.
The parameters `loc`, `scale`, and `concentration` must be shaped in a way
that supports broadcasting (e.g. `loc + scale` + `concentration` is valid).
Args:
loc: Floating point tensor, the location parameter of the distribution(s).
scale: Floating point tensor, the scales of the distribution(s).
scale must contain only positive values.
concentration: Floating point tensor, the concentration of
the distribution(s).
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: `'GeneralizedExtremeValue'`.
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, concentration],
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)
concentration = tensor_util.convert_nonref_to_tensor(
concentration, name='concentration', dtype=dtype)
dtype_util.assert_same_float_dtype([loc, scale, concentration])
# Positive scale is asserted by the incorporated GEV bijector.
self._gev_bijector = gev_cdf_bijector.GeneralizedExtremeValueCDF(
loc=loc, scale=scale, concentration=concentration,
validate_args=validate_args)
batch_shape = distribution_util.get_broadcast_shape(loc, scale,
concentration)
# 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(GeneralizedExtremeValue, self).__init__(
# TODO(b/137665504): Use batch-adding meta-distribution to set the
# batch shape instead of tf.ones.
distribution=uniform.Uniform(
low=np.finfo(dtype_util.as_numpy_dtype(dtype)).tiny,
high=tf.ones(batch_shape, dtype=dtype),
allow_nan_stats=allow_nan_stats),
# The GEV bijector encodes the CDF function as the forward,
# and hence needs to be inverted.
bijector=invert_bijector.Invert(
self._gev_bijector, validate_args=validate_args),
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, values=[loc, scale, low, high]) 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)
tf.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,
graph_parents=[loc, scale, low, high],
name=name)
def precompute_regression_model(
kernel,
observation_index_points,
observations,
index_points=None,
observation_noise_variance=0.,
predictive_noise_variance=None,
mean_fn=None,
jitter=1e-6,
validate_args=False,
allow_nan_stats=False,
name='PrecomputedGaussianProcessRegressionModel'):
"""Returns a GaussianProcessRegressionModel with precomputed quantities.
This differs from the constructor by precomputing quantities associated with
observations in a non-tape safe way. `index_points` is the only parameter
that is allowed to vary (i.e. is a `Variable` / changes after
initialization).
Specifically:
* We make `observation_index_points` and `observations` mandatory
parameters.
* We precompute `kernel(observation_index_points, observation_index_points)`
along with any other associated quantities relating to the `kernel`,
`observations` and `observation_index_points`.
A typical usecase would be optimizing kernel hyperparameters for a
`GaussianProcess`, and computing the posterior predictive with respect to
those optimized hyperparameters and observation / index-points pairs.
WARNING: This method assumes `index_points` is the only varying parameter
(i.e. is a `Variable` / changes after initialization) and hence is not
tape-safe.
Args:
kernel: `PositiveSemidefiniteKernel`-like instance representing the
GP's covariance function.
observation_index_points: `float` `Tensor` representing finite collection,
or batch of collections, of points in the index set for which some data
has been observed. Shape has the form `[b1, ..., bB, e, f1, ..., fF]`
where `F` is the number of feature dimensions and must equal
`kernel.feature_ndims`, and `e` is the number (size) of index points in
each batch. `[b1, ..., bB, e]` must be broadcastable with the shape of
`observations`, and `[b1, ..., bB]` must be broadcastable with the
shapes of all other batched parameters (`kernel.batch_shape`,
`index_points`, etc). The default value is `None`, which corresponds to
the empty set of observations, and simply results in the prior
predictive model (a GP with noise of variance
`predictive_noise_variance`).
observations: `float` `Tensor` representing collection, or batch of
collections, of observations corresponding to
`observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
must be brodcastable with the batch and example shapes of
`observation_index_points`. The batch shape `[b1, ..., bB]` must be
broadcastable with the shapes of all other batched parameters
(`kernel.batch_shape`, `index_points`, etc.). The default value is
`None`, which corresponds to the empty set of observations, and simply
results in the prior predictive model (a GP with noise of variance
`predictive_noise_variance`).
index_points: `float` `Tensor` representing finite collection, or batch of
collections, of points in the index set over which the GP is defined.
Shape has the form `[b1, ..., bB, e, f1, ..., fF]` where `F` is the
number of feature dimensions and must equal `kernel.feature_ndims` and
`e` is the number (size) of index points in each batch. Ultimately this
distribution corresponds to an `e`-dimensional multivariate normal. The
batch shape must be broadcastable with `kernel.batch_shape` and any
batch dims yielded by `mean_fn`.
observation_noise_variance: `float` `Tensor` representing the variance
of the noise in the Normal likelihood distribution of the model. May be
batched, in which case the batch shape must be broadcastable with the
shapes of all other batched parameters (`kernel.batch_shape`,
`index_points`, etc.).
Default value: `0.`
predictive_noise_variance: `float` `Tensor` representing the variance in
the posterior predictive model. If `None`, we simply re-use
`observation_noise_variance` for the posterior predictive noise. If set
explicitly, however, we use this value. This allows us, for example, to
omit predictive noise variance (by setting this to zero) to obtain
noiseless posterior predictions of function values, conditioned on noisy
observations.
mean_fn: Python `callable` that acts on `index_points` to produce a
collection, or batch of collections, of mean values at `index_points`.
Takes a `Tensor` of shape `[b1, ..., bB, f1, ..., fF]` and returns a
`Tensor` whose shape is broadcastable with `[b1, ..., bB]`.
Default value: `None` implies the constant zero function.
jitter: `float` scalar `Tensor` added to the diagonal of the covariance
matrix to ensure positive definiteness of the covariance matrix.
Default value: `1e-6`.
validate_args: Python `bool`, default `False`. When `True` distribution
parameters are checked for validity despite possibly degrading runtime
performance. When `False` invalid inputs may silently render incorrect
outputs.
Default value: `False`.
allow_nan_stats: Python `bool`, default `True`. When `True`,
statistics (e.g., mean, mode, variance) use the value `NaN` to
indicate the result is undefined. When `False`, an exception is raised
if one or more of the statistic's batch members are undefined.
Default value: `False`.
name: Python `str` name prefixed to Ops created by this class.
Default value: 'PrecomputedGaussianProcessRegressionModel'.
Returns
An instance of `GaussianProcessRegressionModel` with precomputed
quantities associated with observations.
"""
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([
index_points, observation_index_points, observations,
observation_noise_variance, predictive_noise_variance, jitter
], tf.float32)
# Convert to tensor arguments that are expected to not be Variables / not
# going to change.
jitter = tf.convert_to_tensor(jitter, dtype=dtype)
observation_index_points = tf.convert_to_tensor(
observation_index_points, dtype=dtype)
observation_noise_variance = tf.convert_to_tensor(
observation_noise_variance, dtype=dtype)
observations = tf.convert_to_tensor(observations, dtype=dtype)
observation_cholesky = kernel.matrix(observation_index_points,
observation_index_points)
broadcast_shape = distribution_util.get_broadcast_shape(
observation_cholesky,
observation_noise_variance[..., tf.newaxis, tf.newaxis])
observation_cholesky = tf.broadcast_to(observation_cholesky,
broadcast_shape)
observation_cholesky = tf.linalg.set_diag(
observation_cholesky,
tf.linalg.diag_part(observation_cholesky) + jitter +
observation_noise_variance[..., tf.newaxis])
observation_cholesky = tf.linalg.cholesky(observation_cholesky)
observation_cholesky_operator = tf.linalg.LinearOperatorLowerTriangular(
observation_cholesky)
conditional_kernel = tfpk.SchurComplement.with_precomputed_divisor(
base_kernel=kernel,
fixed_inputs=observation_index_points,
diag_shift=observation_noise_variance + jitter)
if mean_fn is None:
mean_fn = lambda x: tf.zeros([1], dtype=dtype)
else:
if not callable(mean_fn):
raise ValueError('`mean_fn` must be a Python callable')
diff = observations - mean_fn(observation_index_points)
solve_on_observation = observation_cholesky_operator.solvevec(
observation_cholesky_operator.solvevec(diff), adjoint=True)
def conditional_mean_fn(x):
k_x_obs = kernel.matrix(x, observation_index_points)
return mean_fn(x) + tf.linalg.matvec(k_x_obs,
solve_on_observation)
gprm = GaussianProcessRegressionModel(
kernel=kernel,
observation_index_points=observation_index_points,
observations=observations,
index_points=index_points,
observation_noise_variance=observation_noise_variance,
predictive_noise_variance=predictive_noise_variance,
jitter=jitter,
_conditional_kernel=conditional_kernel,
_conditional_mean_fn=conditional_mean_fn,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
name=name)
return gprm
def test_all_static_shapes_work(self):
x = tf.ones((2, 1, 3))
y = tf.ones((1, 5, 3))
z = tf.ones(())
self.assertAllEqual([2, 5, 3],
distribution_util.get_broadcast_shape(x, y, z))