def gaussian_process_regression_models(draw,
kernel_name=None,
batch_shape=None,
event_dim=None,
feature_dim=None,
feature_ndims=None,
enable_vars=False):
# First draw a kernel.
k, _ = draw(
kernel_hps.base_kernels(
kernel_name=kernel_name,
batch_shape=batch_shape,
event_dim=event_dim,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
# Disable variables
enable_vars=False))
compatible_batch_shape = draw(
tfp_hps.broadcast_compatible_shape(k.batch_shape))
index_points = draw(
kernel_hps.kernel_input(batch_shape=compatible_batch_shape,
example_ndims=1,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
enable_vars=enable_vars,
name='index_points'))
observation_index_points = draw(
kernel_hps.kernel_input(batch_shape=compatible_batch_shape,
example_ndims=1,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
enable_vars=enable_vars,
name='observation_index_points'))
observations = draw(
kernel_hps.kernel_input(
batch_shape=compatible_batch_shape,
example_ndims=1,
# This is the example dimension suggested observation_index_points.
example_dim=int(
observation_index_points.shape[-(feature_ndims + 1)]),
# No feature dimensions.
feature_dim=0,
feature_ndims=0,
enable_vars=enable_vars,
name='observations'))
params = draw(
broadcasting_params('GaussianProcessRegressionModel',
compatible_batch_shape,
event_dim=event_dim,
enable_vars=enable_vars))
gp = tfd.GaussianProcessRegressionModel(
kernel=k,
index_points=index_points,
observation_index_points=observation_index_points,
observations=observations,
observation_noise_variance=params['observation_noise_variance'])
return gp
def gaussian_processes(draw,
kernel_name=None,
batch_shape=None,
event_dim=None,
feature_dim=None,
feature_ndims=None,
enable_vars=False):
# First draw a kernel.
k, _ = draw(kernel_hps.base_kernels(
kernel_name=kernel_name,
batch_shape=batch_shape,
event_dim=event_dim,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
# Disable variables
enable_vars=False))
compatible_batch_shape = draw(
tfp_hps.broadcast_compatible_shape(k.batch_shape))
index_points = draw(kernel_hps.kernel_input(
batch_shape=compatible_batch_shape,
example_ndims=1,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
enable_vars=enable_vars,
name='index_points'))
params = draw(broadcasting_params(
'GaussianProcess',
compatible_batch_shape,
event_dim=event_dim,
enable_vars=enable_vars))
gp = tfd.GaussianProcess(
kernel=k, index_points=index_points,
observation_noise_variance=params['observation_noise_variance'])
return gp
def student_t_processes(draw,
kernel_name=None,
batch_shape=None,
event_dim=None,
feature_dim=None,
feature_ndims=None,
enable_vars=False):
# First draw a kernel.
k, _ = draw(
kernel_hps.base_kernels(
kernel_name=kernel_name,
batch_shape=batch_shape,
event_dim=event_dim,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
# Disable variables
enable_vars=False))
compatible_batch_shape = draw(
tfp_hps.broadcast_compatible_shape(k.batch_shape))
index_points = draw(
kernel_hps.kernel_input(batch_shape=compatible_batch_shape,
example_ndims=1,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
enable_vars=enable_vars,
name='index_points'))
params = draw(
broadcasting_params('StudentTProcess',
compatible_batch_shape,
event_dim=event_dim,
enable_vars=enable_vars))
stp = tfd.StudentTProcess(
kernel=k,
index_points=index_points,
# The Student-T Process can encounter cholesky decomposition errors,
# so use a large jitter to avoid that.
jitter=1e-1,
df=params['df'])
return stp
def student_t_processes(draw,
kernel_name=None,
batch_shape=None,
event_dim=None,
feature_dim=None,
feature_ndims=None,
enable_vars=False):
# First draw a kernel.
k, _ = draw(
kernel_hps.base_kernels(
kernel_name=kernel_name,
batch_shape=batch_shape,
event_dim=event_dim,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
# Disable variables
enable_vars=False))
compatible_batch_shape = draw(
tfp_hps.broadcast_compatible_shape(k.batch_shape))
index_points = draw(
kernel_hps.kernel_input(batch_shape=compatible_batch_shape,
example_ndims=1,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
enable_vars=enable_vars,
name='index_points'))
params = draw(
broadcasting_params('StudentTProcess',
compatible_batch_shape,
event_dim=event_dim,
enable_vars=enable_vars))
stp = tfd.StudentTProcess(
kernel=k,
index_points=index_points,
cholesky_fn=lambda x: marginal_fns.retrying_cholesky(x)[0],
df=params['df'],
observation_noise_variance=params['observation_noise_variance'])
return stp
def student_t_process_regression_models(draw,
kernel_name=None,
batch_shape=None,
event_dim=None,
feature_dim=None,
feature_ndims=None,
enable_vars=False):
# First draw a kernel.
k, _ = draw(kernel_hps.base_kernels(
kernel_name=kernel_name,
batch_shape=batch_shape,
event_dim=event_dim,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
# Disable variables
enable_vars=False))
compatible_batch_shape = draw(
tfp_hps.broadcast_compatible_shape(k.batch_shape))
index_points = draw(kernel_hps.kernel_input(
batch_shape=compatible_batch_shape,
example_ndims=1,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
enable_vars=enable_vars,
name='index_points'))
hp.note('Index points:\n{}'.format(repr(index_points)))
observation_index_points = draw(
kernel_hps.kernel_input(
batch_shape=compatible_batch_shape,
example_ndims=1,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
enable_vars=enable_vars,
name='observation_index_points'))
hp.note('Observation index points:\n{}'.format(
repr(observation_index_points)))
observations = draw(kernel_hps.kernel_input(
batch_shape=compatible_batch_shape,
example_ndims=1,
# This is the example dimension suggested observation_index_points.
example_dim=int(observation_index_points.shape[-(feature_ndims + 1)]),
# No feature dimensions.
feature_dim=0,
feature_ndims=0,
enable_vars=enable_vars,
name='observations'))
hp.note('Observations:\n{}'.format(repr(observations)))
params = draw(broadcasting_params(
'StudentTProcessRegressionModel',
compatible_batch_shape,
event_dim=event_dim,
enable_vars=enable_vars))
hp.note('Params:\n{}'.format(repr(params)))
stp = tfd.StudentTProcessRegressionModel(
# Ensure that the `df` parameter is not a `Variable` since we pass
# in a `DeferredTensor` of the `df` parameter.
df=tf.convert_to_tensor(params['df']),
kernel=k,
index_points=index_points,
observation_index_points=observation_index_points,
observations=observations,
cholesky_fn=lambda x: marginal_fns.retrying_cholesky(x)[0],
observation_noise_variance=params['observation_noise_variance'])
return stp
def variational_gaussian_processes(draw,
kernel_name=None,
batch_shape=None,
event_dim=None,
feature_dim=None,
feature_ndims=None,
enable_vars=False):
# First draw a kernel.
k, _ = draw(kernel_hps.base_kernels(
kernel_name=kernel_name,
batch_shape=batch_shape,
event_dim=event_dim,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
# Disable variables
enable_vars=False))
compatible_batch_shape = draw(
tfp_hps.broadcast_compatible_shape(k.batch_shape))
index_points = draw(kernel_hps.kernel_input(
batch_shape=compatible_batch_shape,
example_ndims=1,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
enable_vars=enable_vars,
name='index_points'))
hp.note('Index points:\n{}'.format(repr(index_points)))
inducing_index_points = draw(kernel_hps.kernel_input(
batch_shape=compatible_batch_shape,
example_ndims=1,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
enable_vars=enable_vars,
name='inducing_index_points'))
hp.note('Inducing index points:\n{}'.format(repr(inducing_index_points)))
num_inducing_points = int(inducing_index_points.shape[-(feature_ndims + 1)])
variational_inducing_observations_loc = draw(kernel_hps.kernel_input(
batch_shape=compatible_batch_shape,
example_ndims=1,
example_dim=num_inducing_points,
feature_dim=0,
feature_ndims=0,
enable_vars=enable_vars,
name='variational_inducing_observations_loc'))
hp.note('Variational inducing observations loc:\n{}'.format(
repr(variational_inducing_observations_loc)))
variational_inducing_observations_scale = draw(tfp_hps.tensors_in_support(
support=tfp_hps.Support.MATRIX_LOWER_TRIL_POSITIVE_DEFINITE,
batch_shape=compatible_batch_shape.as_list(),
event_dim=num_inducing_points,
dtype=np.float64))
hp.note('Variational inducing observations scale:\n{}'.format(
repr(variational_inducing_observations_scale)))
params = draw(broadcasting_params(
'GaussianProcessRegressionModel',
compatible_batch_shape,
event_dim=event_dim,
enable_vars=enable_vars))
hp.note('Params:\n{}'.format(repr(params)))
vgp = tfd.VariationalGaussianProcess(
kernel=k,
index_points=index_points,
inducing_index_points=inducing_index_points,
variational_inducing_observations_loc=(
variational_inducing_observations_loc),
variational_inducing_observations_scale=(
variational_inducing_observations_scale),
observation_noise_variance=params[
'observation_noise_variance'])
return vgp
