def scalar_constrainer(support):
def nonzero(x):
return tf.where(tf.equal(x, 0), 1e-6, x)
constrainers = {
Support.SCALAR_IN_0_1: tf.math.sigmoid,
Support.SCALAR_GT_NEG1: tfp_hps.softplus_plus_eps(-1 + 1e-6),
Support.SCALAR_NON_ZERO: nonzero,
Support.SCALAR_IN_NEG1_1: lambda x: tf.math.tanh(x) * (1 - 1e-6),
Support.SCALAR_NON_NEGATIVE: tf.math.softplus,
Support.SCALAR_POSITIVE: tfp_hps.softplus_plus_eps(),
Support.SCALAR_UNCONSTRAINED: tf.identity,
}
return constrainers[support]
def generalized_paretos(draw, batch_shape=None):
if batch_shape is None:
batch_shape = draw(tfp_hps.shapes())
constraints = dict(
loc=tfp_hps.identity_fn,
scale=tfp_hps.softplus_plus_eps(),
concentration=lambda x: tf.math.tanh(x) * 0.24) # <.25==safe for variance
params = draw(
tfp_hps.broadcasting_params(
batch_shape,
params_event_ndims=dict(loc=0, scale=0, concentration=0),
constraint_fn_for=constraints.get))
dist = tfd.GeneralizedPareto(validate_args=draw(hps.booleans()), **params)
if dist.batch_shape != batch_shape:
raise AssertionError('batch_shape mismatch: expect {} but got {}'.format(
batch_shape, dist))
return dist
# Ensure that low < high.
return dict(d, high=tfp_hps.ensure_high_gt_low(d['low'], d['high']))
def fix_softclip(d):
# Ensure that low < high.
return dict(d, high=tfp_hps.ensure_high_gt_low(d['low'], d['high']))
def fix_rational_quadratic(d):
return dict(d, range_min=-1)
CONSTRAINTS = {
'concentration':
tfp_hps.softplus_plus_eps(),
'concentration0':
tfp_hps.softplus_plus_eps(),
'concentration1':
tfp_hps.softplus_plus_eps(),
'hinge_softness':
tfp_hps.softplus_plus_eps(),
'power':
# Restrict to positive since `Invert(Power(...))` tests the negation.
tfp_hps.softplus_plus_eps(),
'rate':
tfp_hps.softplus_plus_eps(),
'scale':
tfp_hps.softplus_plus_eps(),
'tailweight':
tfp_hps.softplus_plus_eps(),
'method `inverse_log_det_jacobian` of {}'.format(bijector),
max_permissible=max_permitted):
tape.watch(wrt_vars)
# TODO(b/73073515): Fix graph mode gradients with bijector caching.
xs = bijector.inverse_log_det_jacobian(ys + 0, event_ndims=event_ndims)
grads = tape.gradient(xs, wrt_vars)
assert_no_none_grad(bijector, 'inverse_log_det_jacobian', wrt_vars, grads)
def ensure_nonzero(x):
return tf.where(x < 1e-6, tf.constant(1e-3, x.dtype), x)
CONSTRAINTS = {
'concentration':
tfp_hps.softplus_plus_eps(),
'concentration0':
tfp_hps.softplus_plus_eps(),
'concentration1':
tfp_hps.softplus_plus_eps(),
'scale':
tfp_hps.softplus_plus_eps(),
'tailweight':
tfp_hps.softplus_plus_eps(),
'AffineScalar.scale':
tfp_hps.softplus_plus_eps(),
'bin_widths':
bijector_hps.spline_bin_size_constraint,
'bin_heights':
bijector_hps.spline_bin_size_constraint,
'knot_slopes':
return dict(d, peak=peak, high=high)
def fix_wishart(d):
df = d['df']
scale = d.get('scale', d.get('scale_tril'))
return dict(d, df=tf.maximum(df, tf.cast(scale.shape[-1], df.dtype)))
CONSTRAINTS = {
'atol':
tf.math.softplus,
'rtol':
tf.math.softplus,
'concentration':
tfp_hps.softplus_plus_eps(),
'GeneralizedPareto.concentration': # Permits +ve and -ve concentrations.
lambda x: tf.math.tanh(x) * 0.24,
'concentration0':
tfp_hps.softplus_plus_eps(),
'concentration1':
tfp_hps.softplus_plus_eps(),
'covariance_matrix':
tfp_hps.positive_definite,
'df':
tfp_hps.softplus_plus_eps(),
'InverseGaussian.loc':
tfp_hps.softplus_plus_eps(),
'VonMisesFisher.mean_direction': # max ndims is 5
lambda x: tf.math.l2_normalize(tf.math.sigmoid(x[..., :5]) + 1e-6, -1),
'Categorical.probs':
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 reconstructing the kernel works
with tfp_hps.no_tf_rank_errors():
diag2 = type(kernel)(**kernel._parameters).apply(
xs, xs, example_ndims=example_ndims)
self.assertAllClose(diag, diag2)
CONSTRAINTS = {
# Keep amplitudes large enough so that the matrices are well conditioned.
'amplitude': tfp_hps.softplus_plus_eps(1.),
'bias_variance': tfp_hps.softplus_plus_eps(1.),
'slope_variance': tfp_hps.softplus_plus_eps(1.),
'exponent': tfp_hps.softplus_plus_eps(),
'length_scale': tfp_hps.softplus_plus_eps(),
'period': tfp_hps.softplus_plus_eps(),
'scale_mixture_rate': tfp_hps.softplus_plus_eps(),
}
def constraint_for(kernel_name=None, param=None):
if param is not None:
return CONSTRAINTS.get('{}.{}'.format(kernel_name, param),
CONSTRAINTS.get(param, tfp_hps.identity_fn))
return CONSTRAINTS.get(kernel_name, tfp_hps.identity_fn)
return dict(d, peak=peak, high=high)
def fix_wishart(d):
df = d['df']
scale = d.get('scale', d.get('scale_tril'))
return dict(d, df=tf.maximum(df, tf.cast(scale.shape[-1], df.dtype)))
CONSTRAINTS = {
'atol':
tf.math.softplus,
'rtol':
tf.math.softplus,
'concentration':
tfp_hps.softplus_plus_eps(),
'GeneralizedPareto.concentration': # Permits +ve and -ve concentrations.
lambda x: tf.math.tanh(x) * 0.24,
'concentration0':
tfp_hps.softplus_plus_eps(),
'concentration1':
tfp_hps.softplus_plus_eps(),
'covariance_matrix':
tfp_hps.positive_definite,
'df':
tfp_hps.softplus_plus_eps(),
'Chi2WithAbsDf.df':
tfp_hps.softplus_plus_eps(1), # does floor(abs(x)) for some reason
'InverseGaussian.loc':
tfp_hps.softplus_plus_eps(),
'VonMisesFisher.mean_direction': # max ndims is 5
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_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 spectral_mixtures(draw,
batch_shape=None,
event_dim=None,
feature_dim=None,
feature_ndims=None,
enable_vars=None,
depth=None):
"""Strategy for drawing `SpectralMixture` 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))
num_mixtures = draw(hps.integers(min_value=2, max_value=5))
logits = draw(
kernel_input(batch_shape=batch_shape,
example_ndims=0,
feature_dim=num_mixtures,
feature_ndims=1))
locs = draw(
kernel_input(batch_shape=batch_shape,
example_ndims=1,
example_dim=num_mixtures,
feature_dim=feature_dim,
feature_ndims=feature_ndims))
scales = tfp_hps.softplus_plus_eps()(draw(
kernel_input(batch_shape=batch_shape,
example_ndims=1,
example_dim=num_mixtures,
feature_dim=feature_dim,
feature_ndims=feature_ndims)))
hp.note(f'Forming SpectralMixture kernel with logits: {logits} '
f'locs: {locs} and scales: {scales}')
spectral_mixture_params = {
'locs': locs,
'logits': logits,
'scales': scales
}
kernel_variable_names = []
for param_name in spectral_mixture_params:
if enable_vars and draw(hps.booleans()):
kernel_variable_names.append(param_name)
spectral_mixture_params[param_name] = tf.Variable(
spectral_mixture_params[param_name], name=param_name)
if draw(hps.booleans()):
spectral_mixture_params[
param_name] = tfp_hps.defer_and_count_usage(
spectral_mixture_params[param_name])
result_kernel = tfpk.SpectralMixture(
logits=spectral_mixture_params['logits'],
locs=spectral_mixture_params['locs'],
scales=spectral_mixture_params['scales'],
feature_ndims=feature_ndims,
validate_args=True)
return result_kernel, kernel_variable_names
def changepoints(draw,
batch_shape=None,
event_dim=None,
feature_dim=None,
feature_ndims=None,
enable_vars=None,
depth=None):
"""Strategy for drawing `Changepoint` 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 `Changepoint` 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))
num_kernels = draw(hps.integers(min_value=2, max_value=4))
inner_kernels = []
kernel_variable_names = []
for _ in range(num_kernels):
base_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))
inner_kernels.append(base_kernel)
kernel_variable_names += variable_names
constraints = dict(
locs=lambda x: tf.cumsum(tf.math.abs(x) + 1e-3, axis=-1),
slopes=tfp_hps.softplus_plus_eps())
params = draw(
tfp_hps.broadcasting_params(batch_shape,
event_dim=num_kernels - 1,
params_event_ndims=dict(locs=1, slopes=1),
constraint_fn_for=constraints.get))
params = {k: tf.cast(params[k], tf.float64) for k in params}
if enable_vars and draw(hps.booleans()):
kernel_variable_names.append('locs')
kernel_variable_names.append('slopes')
params['locs'] = tf.Variable(params['locs'], name='locs')
params['slopes'] = tf.Variable(params['slopes'], name='slopes')
result_kernel = tfpk.ChangePoint(kernels=inner_kernels,
validate_args=True,
**params)
return result_kernel, kernel_variable_names
max_permitted = (2 if hasattr(
bijector.bijector, '_forward_log_det_jacobian') else 4)
with tfp_hps.assert_no_excessive_var_usage(
'method `inverse_log_det_jacobian` of {}'.format(bijector),
max_permissible=max_permitted):
tape.watch(wrt_vars)
# TODO(b/73073515): Fix graph mode gradients with bijector caching.
xs = bijector.inverse_log_det_jacobian(ys + 0,
event_ndims=event_ndims)
grads = tape.gradient(xs, wrt_vars)
assert_no_none_grad(bijector, 'inverse_log_det_jacobian', wrt_vars,
grads)
CONSTRAINTS = {
'concentration0': tfp_hps.softplus_plus_eps(),
'concentration1': tfp_hps.softplus_plus_eps(),
'scale': tfp_hps.softplus_plus_eps(),
'tailweight': tfp_hps.softplus_plus_eps(),
'AffineScalar.scale': tfp_hps.softplus_plus_eps(),
}
def constraint_for(bijector_name=None, param=None):
if param is not None:
return CONSTRAINTS.get('{}.{}'.format(bijector_name, param),
CONSTRAINTS.get(param, tfp_hps.identity_fn))
return CONSTRAINTS.get(bijector_name, tfp_hps.identity_fn)
if __name__ == '__main__':