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 testCompositeTensor(self, kernel_name, data):
kernel, _ = data.draw(
kernel_hps.kernels(kernel_name=kernel_name,
event_dim=2,
feature_dim=2,
feature_ndims=1,
enable_vars=True))
self.assertIsInstance(kernel, tf.__internal__.CompositeTensor)
xs = tf.identity(
data.draw(
kernel_hps.kernel_input(batch_shape=[],
example_ndims=1,
feature_dim=2,
feature_ndims=1)))
with tfp_hps.no_tf_rank_errors():
diag = kernel.apply(xs, xs, example_ndims=1)
# Test flatten/unflatten.
flat = tf.nest.flatten(kernel, expand_composites=True)
unflat = tf.nest.pack_sequence_as(kernel, flat, expand_composites=True)
# Test tf.function.
@tf.function
def diag_fn(k):
return k.apply(xs, xs, example_ndims=1)
self.evaluate([v.initializer for v in kernel.variables])
with tfp_hps.no_tf_rank_errors():
self.assertAllClose(diag, diag_fn(kernel))
self.assertAllClose(diag, diag_fn(unflat))
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 testKernelGradient(self, kernel_name, data):
event_dim = data.draw(hps.integers(min_value=2, max_value=3))
feature_ndims = data.draw(hps.integers(min_value=1, max_value=2))
feature_dim = data.draw(hps.integers(min_value=2, max_value=4))
batch_shape = data.draw(tfp_hps.shapes(max_ndims=2))
kernel, kernel_parameter_variable_names = data.draw(
kernel_hps.kernels(batch_shape=batch_shape,
kernel_name=kernel_name,
event_dim=event_dim,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
enable_vars=True))
# Check that variable parameters get passed to the kernel.variables
kernel_variables_names = [
v.name.strip('_0123456789:') for v in kernel.variables
]
kernel_parameter_variable_names = [
n.strip('_0123456789:') for n in kernel_parameter_variable_names
]
self.assertEqual(set(kernel_parameter_variable_names),
set(kernel_variables_names))
example_ndims = data.draw(hps.integers(min_value=1, max_value=2))
input_batch_shape = data.draw(
tfp_hps.broadcast_compatible_shape(kernel.batch_shape))
xs = tf.identity(
data.draw(
kernel_hps.kernel_input(batch_shape=input_batch_shape,
example_ndims=example_ndims,
feature_dim=feature_dim,
feature_ndims=feature_ndims)))
# Check that we pick up all relevant kernel parameters.
wrt_vars = [xs] + list(kernel.variables)
self.evaluate([v.initializer for v in kernel.variables])
max_permissible = 2 + EXTRA_TENSOR_CONVERSION_KERNELS.get(
kernel_name, 0)
with tf.GradientTape() as tape:
with tfp_hps.assert_no_excessive_var_usage(
'method `apply` of {}'.format(kernel),
max_permissible=max_permissible):
tape.watch(wrt_vars)
with tfp_hps.no_tf_rank_errors():
diag = kernel.apply(xs, xs, example_ndims=example_ndims)
grads = tape.gradient(diag, wrt_vars)
assert_no_none_grad(kernel, 'apply', wrt_vars, grads)
# Check that copying the kernel works.
with tfp_hps.no_tf_rank_errors():
diag2 = self.evaluate(kernel.copy().apply(
xs, xs, example_ndims=example_ndims))
self.assertAllClose(diag, diag2)
def testKernelGradient(self, kernel_name, data):
if tf.executing_eagerly() != (FLAGS.tf_mode == 'eager'):
return
event_dim = data.draw(hps.integers(min_value=2, max_value=6))
feature_ndims = data.draw(hps.integers(min_value=1, max_value=4))
feature_dim = data.draw(hps.integers(min_value=2, max_value=6))
kernel, kernel_parameter_variable_names = data.draw(
kernel_hps.kernels(
kernel_name=kernel_name,
event_dim=event_dim,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
enable_vars=True))
# Check that variable parameters get passed to the kernel.variables
kernel_variables_names = [
v.name.strip('_0123456789:') for v in kernel.variables]
self.assertEqual(
set(kernel_parameter_variable_names),
set(kernel_variables_names))
example_ndims = data.draw(hps.integers(min_value=1, max_value=3))
input_batch_shape = data.draw(tfp_hps.broadcast_compatible_shape(
kernel.batch_shape))
xs = tf.identity(data.draw(kernel_hps.kernel_input(
batch_shape=input_batch_shape,
example_ndims=example_ndims,
feature_dim=feature_dim,
feature_ndims=feature_ndims)))
# Check that we pick up all relevant kernel parameters.
wrt_vars = [xs] + list(kernel.variables)
self.evaluate([v.initializer for v in kernel.variables])
with tf.GradientTape() as tape:
with tfp_hps.assert_no_excessive_var_usage(
'method `apply` of {}'.format(kernel)):
tape.watch(wrt_vars)
diag = kernel.apply(xs, xs, example_ndims=example_ndims)
grads = tape.gradient(diag, wrt_vars)
assert_no_none_grad(kernel, 'apply', wrt_vars, grads)
self.assertAllClose(
diag,
type(kernel)(**kernel._parameters).apply(
xs, xs, example_ndims=example_ndims))
def _test_slicing(self, data, kernel_name, kernel, feature_dim,
feature_ndims):
example_ndims = data.draw(hps.integers(min_value=0, max_value=2))
batch_shape = kernel.batch_shape
slices = data.draw(tfp_hps.valid_slices(batch_shape))
slice_str = 'kernel[{}]'.format(', '.join(
tfp_hps.stringify_slices(slices)))
# Make sure the slice string appears in Hypothesis' attempted example log
hp.note('Using slice ' + slice_str)
if not slices: # Nothing further to check.
return
sliced_zeros = np.zeros(batch_shape)[slices]
sliced_kernel = kernel[slices]
hp.note('Using sliced kernel {}.'.format(sliced_kernel))
hp.note('Using sliced zeros {}.'.format(sliced_zeros.shape))
# Check that slicing modifies batch shape as expected.
self.assertAllEqual(sliced_zeros.shape, sliced_kernel.batch_shape)
xs = tf.identity(
data.draw(
kernel_hps.kernel_input(batch_shape=[],
example_ndims=example_ndims,
feature_dim=feature_dim,
feature_ndims=feature_ndims)))
# Check that apply of sliced kernels executes.
with tfp_hps.no_tf_rank_errors():
results = self.evaluate(
kernel.apply(xs, xs, example_ndims=example_ndims))
hp.note('Using results shape {}.'.format(results.shape))
sliced_results = self.evaluate(
sliced_kernel.apply(xs, xs, example_ndims=example_ndims))
# Come up with the slices for apply (which must also include example dims).
apply_slices = (tuple(slices) if isinstance(
slices, collections.Sequence) else (slices, ))
apply_slices += tuple([slice(None)] * example_ndims)
# Check that sampling a sliced kernel produces the same shape as
# slicing the samples from the original.
self.assertAllClose(results[apply_slices], sliced_results)
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 schur_complements(draw,
batch_shape=None,
event_dim=None,
feature_dim=None,
feature_ndims=None,
enable_vars=None,
depth=None):
"""Strategy for drawing `SchurComplement` kernels.
The underlying kernel is drawn from the `kernels` strategy.
Args:
draw: Hypothesis strategy sampler supplied by `@hps.composite`.
batch_shape: An optional `TensorShape`. The batch shape of the resulting
Kernel. Hypothesis will pick a batch shape if omitted.
event_dim: Optional Python int giving the size of each of the
kernel's parameters' event dimensions. This is shared across all
parameters, permitting square event matrices, compatible location and
scale Tensors, etc. If omitted, Hypothesis will choose one.
feature_dim: Optional Python int giving the size of each feature dimension.
If omitted, Hypothesis will choose one.
feature_ndims: Optional Python int stating the number of feature dimensions
inputs will have. If omitted, Hypothesis will choose one.
enable_vars: TODO(bjp): Make this `True` all the time and put variable
initialization in slicing_test. If `False`, the returned parameters are
all Tensors, never Variables or DeferredTensor.
depth: Python `int` giving maximum nesting depth of compound kernel.
Returns:
kernels: A strategy for drawing `SchurComplement` kernels with the specified
`batch_shape` (or an arbitrary one if omitted).
"""
if depth is None:
depth = draw(depths())
if batch_shape is None:
batch_shape = draw(tfp_hps.shapes())
if event_dim is None:
event_dim = draw(hps.integers(min_value=2, max_value=6))
if feature_dim is None:
feature_dim = draw(hps.integers(min_value=2, max_value=6))
if feature_ndims is None:
feature_ndims = draw(hps.integers(min_value=2, max_value=6))
base_kernel, kernel_variable_names = draw(
kernels(batch_shape=batch_shape,
event_dim=event_dim,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
enable_vars=False,
depth=depth - 1))
# SchurComplement requires the inputs to have one example dimension.
fixed_inputs = draw(
kernel_hps.kernel_input(batch_shape=batch_shape,
example_ndims=1,
feature_dim=feature_dim,
feature_ndims=feature_ndims))
# Positive shift to ensure the divisor matrix is PD.
diag_shift = draw(
hpnp.arrays(dtype=np.float32,
shape=batch_shape,
elements=hps.floats(1,
100,
allow_nan=False,
allow_infinity=False)))
hp.note('Forming SchurComplement kernel with fixed_inputs: {} '
'and diag_shift: {}'.format(fixed_inputs, diag_shift))
schur_complement_params = {
'fixed_inputs': fixed_inputs,
'diag_shift': diag_shift
}
for param_name in schur_complement_params:
if enable_vars and draw(hps.booleans()):
kernel_variable_names.append(param_name)
schur_complement_params[param_name] = tf.Variable(
schur_complement_params[param_name], name=param_name)
if draw(hps.booleans()):
schur_complement_params[
param_name] = tfp_hps.defer_and_count_usage(
schur_complement_params[param_name])
result_kernel = tfp.positive_semidefinite_kernels.SchurComplement(
base_kernel=base_kernel,
fixed_inputs=schur_complement_params['fixed_inputs'],
diag_shift=schur_complement_params['diag_shift'],
validate_args=True)
return result_kernel, kernel_variable_names
def feature_scaleds(draw,
batch_shape=None,
event_dim=None,
feature_dim=None,
feature_ndims=None,
enable_vars=None,
depth=None):
"""Strategy for drawing `FeatureScaled` kernels.
The underlying kernel is drawn from the `kernels` strategy.
Args:
draw: Hypothesis strategy sampler supplied by `@hps.composite`.
batch_shape: An optional `TensorShape`. The batch shape of the resulting
Kernel. Hypothesis will pick a batch shape if omitted.
event_dim: Optional Python int giving the size of each of the
kernel's parameters' event dimensions. This is shared across all
parameters, permitting square event matrices, compatible location and
scale Tensors, etc. If omitted, Hypothesis will choose one.
feature_dim: Optional Python int giving the size of each feature dimension.
If omitted, Hypothesis will choose one.
feature_ndims: Optional Python int stating the number of feature dimensions
inputs will have. If omitted, Hypothesis will choose one.
enable_vars: TODO(bjp): Make this `True` all the time and put variable
initialization in slicing_test. If `False`, the returned parameters are
all Tensors, never Variables or DeferredTensor.
depth: Python `int` giving maximum nesting depth of compound kernel.
Returns:
kernels: A strategy for drawing `FeatureScaled` kernels with the specified
`batch_shape` (or an arbitrary one if omitted).
"""
if depth is None:
depth = draw(depths())
if batch_shape is None:
batch_shape = draw(tfp_hps.shapes())
if event_dim is None:
event_dim = draw(hps.integers(min_value=2, max_value=6))
if feature_dim is None:
feature_dim = draw(hps.integers(min_value=2, max_value=6))
if feature_ndims is None:
feature_ndims = draw(hps.integers(min_value=2, max_value=6))
base_kernel, kernel_variable_names = draw(
kernels(batch_shape=batch_shape,
event_dim=event_dim,
feature_dim=feature_dim,
feature_ndims=feature_ndims,
enable_vars=False,
depth=depth - 1))
scale_diag = tfp_hps.softplus_plus_eps()(draw(
kernel_hps.kernel_input(batch_shape=batch_shape,
example_ndims=0,
feature_dim=feature_dim,
feature_ndims=feature_ndims)))
hp.note(
'Forming FeatureScaled kernel with scale_diag: {} '.format(scale_diag))
if enable_vars and draw(hps.booleans()):
kernel_variable_names.append('scale_diag')
scale_diag = tf.Variable(scale_diag, name='scale_diag')
# Don't enable variable counting. This is because rescaling is
# done for each input, which will exceed two convert_to_tensor calls.
result_kernel = tfp.positive_semidefinite_kernels.FeatureScaled(
kernel=base_kernel, scale_diag=scale_diag, validate_args=True)
return result_kernel, kernel_variable_names
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
