def _inverse(self, y):
map_values = tf.convert_to_tensor(self.map_values)
flat_y = tf.reshape(y, shape=[-1])
# Search for the indices of map_values that are closest to flat_y.
# Since map_values is strictly increasing, the closest is either the
# first one that is strictly greater than flat_y, or the one before it.
upper_candidates = tf.minimum(
tf.size(map_values) - 1,
tf.searchsorted(map_values, values=flat_y, side='right'))
lower_candidates = tf.maximum(0, upper_candidates - 1)
candidates = tf.stack([lower_candidates, upper_candidates], axis=-1)
lower_cand_diff = tf.abs(flat_y - self._forward(lower_candidates))
upper_cand_diff = tf.abs(flat_y - self._forward(upper_candidates))
if self.validate_args:
with tf.control_dependencies([
assert_util.assert_near(tf.minimum(lower_cand_diff,
upper_cand_diff),
0,
message='inverse value not found')
]):
candidates = tf.identity(candidates)
candidate_selector = tf.stack([
tf.range(tf.size(flat_y), dtype=tf.int32),
tf.argmin([lower_cand_diff, upper_cand_diff], output_type=tf.int32)
],
axis=-1)
return tf.reshape(tf.gather_nd(candidates, candidate_selector),
shape=y.shape)
def _parameter_control_dependencies(self, is_init):
if not self.validate_args:
return []
mean_direction = tf.convert_to_tensor(self.mean_direction)
concentration = tf.convert_to_tensor(self.concentration)
assertions = []
if is_init != tensor_util.is_ref(self._mean_direction):
assertions.append(
assert_util.assert_greater(
tf.shape(mean_direction)[-1],
1,
message='`mean_direction` may not have scalar event shape'))
assertions.append(
assert_util.assert_less_equal(
tf.shape(mean_direction)[-1],
5,
message='von Mises-Fisher ndims > 5 is not currently supported'))
assertions.append(
assert_util.assert_near(
1.,
tf.linalg.norm(mean_direction, axis=-1),
message='`mean_direction` must be unit-length'))
if is_init != tensor_util.is_ref(self._concentration):
assertions.append(
assert_util.assert_non_negative(
concentration, message='`concentration` must be non-negative'))
return assertions
def _sample_control_dependencies(self, x):
assertions = []
if tensorshape_util.is_fully_defined(x.shape[-2:]):
if not (tensorshape_util.dims(x.shape)[-2] ==
tensorshape_util.dims(x.shape)[-1] == self.dimension):
raise ValueError(
'Input dimension mismatch: expected [..., {}, {}], got {}'.
format(self.dimension, self.dimension,
tensorshape_util.dims(x.shape)))
elif self.validate_args:
msg = 'Input dimension mismatch: expected [..., {}, {}], got {}'.format(
self.dimension, self.dimension, tf.shape(x))
assertions.append(
assert_util.assert_equal(tf.shape(x)[-2],
self.dimension,
message=msg))
assertions.append(
assert_util.assert_equal(tf.shape(x)[-1],
self.dimension,
message=msg))
if self.validate_args:
assertions.append(
assert_util.assert_near(
x,
tf.linalg.band_part(x, -1, 0),
message='Cholesky factors must be lower triangular.'))
return assertions
def _parameter_control_dependencies(self, is_init):
assertions = []
logits = self._logits
probs = self._probs
param, name = (probs, 'probs') if logits is None else (logits, 'logits')
# In init, we can always build shape and dtype checks because
# we assume shape doesn't change for Variable backed args.
if is_init:
if not dtype_util.is_floating(param.dtype):
raise TypeError('Argument `{}` must having floating type.'.format(name))
msg = 'Argument `{}` must have rank at least 1.'.format(name)
shape_static = tensorshape_util.dims(param.shape)
if shape_static is not None:
if len(shape_static) < 1:
raise ValueError(msg)
elif self.validate_args:
param = tf.convert_to_tensor(param)
assertions.append(
assert_util.assert_rank_at_least(param, 1, message=msg))
with tf.control_dependencies(assertions):
param = tf.identity(param)
msg1 = 'Argument `{}` must have final dimension >= 1.'.format(name)
msg2 = 'Argument `{}` must have final dimension <= {}.'.format(
name, dtype_util.max(tf.int32))
event_size = shape_static[-1] if shape_static is not None else None
if event_size is not None:
if event_size < 1:
raise ValueError(msg1)
if event_size > dtype_util.max(tf.int32):
raise ValueError(msg2)
elif self.validate_args:
param = tf.convert_to_tensor(param)
assertions.append(assert_util.assert_greater_equal(
tf.shape(param)[-1], 1, message=msg1))
# NOTE: For now, we leave out a runtime assertion that
# `tf.shape(param)[-1] <= tf.int32.max`. An earlier `tf.shape` call
# will fail before we get to this point.
if not self.validate_args:
assert not assertions # Should never happen.
return []
if probs is not None:
probs = param # reuse tensor conversion from above
if is_init != tensor_util.is_ref(probs):
probs = tf.convert_to_tensor(probs)
one = tf.ones([], dtype=probs.dtype)
assertions.extend([
assert_util.assert_non_negative(probs),
assert_util.assert_less_equal(probs, one),
assert_util.assert_near(
tf.reduce_sum(probs, axis=-1), one,
message='Argument `probs` must sum to 1.'),
])
return assertions
def _inverse(self, y):
# To derive the inverse mapping note that:
# y[i] = exp(x[i]) / normalization
# and
# y[end] = 1 / normalization.
# Thus:
# x[i] = log(exp(x[i])) - log(y[end]) - log(normalization)
# = log(exp(x[i])/normalization) - log(y[end])
# = log(y[i]) - log(y[end])
# Do this first to make sure CSE catches that it'll happen again in
# _inverse_log_det_jacobian.
assertions = []
if self.validate_args:
assertions.append(assert_util.assert_near(
tf.reduce_sum(y, axis=-1),
tf.ones([], y.dtype),
2. * np.finfo(dtype_util.as_numpy_dtype(y.dtype)).eps,
message='Last dimension of `y` must sum to `1`.'))
assertions.append(assert_util.assert_less_equal(
y, tf.ones([], y.dtype),
message='Elements of `y` must be less than or equal to `1`.'))
assertions.append(assert_util.assert_non_negative(
y, message='Elements of `y` must be non-negative.'))
with tf.control_dependencies(assertions):
x = tf.math.log(y)
x, log_normalization = tf.split(x, num_or_size_splits=[-1, 1], axis=-1)
return x - log_normalization
def _parameter_control_dependencies(self, is_init):
if not self.validate_args:
return []
mean_direction = tf.convert_to_tensor(self.mean_direction)
concentration = tf.convert_to_tensor(self.concentration)
assertions = []
if is_init != tensor_util.is_ref(self._mean_direction):
assertions.append(
assert_util.assert_greater(
tf.shape(mean_direction)[-1],
1,
message=
'`mean_direction` must be a vector of at least size 2.'))
assertions.append(
assert_util.assert_near(
tf.cast(1., self.dtype),
tf.linalg.norm(mean_direction, axis=-1),
message='`mean_direction` must be unit-length'))
if is_init != tensor_util.is_ref(self._concentration):
assertions.append(
assert_util.assert_non_negative(
concentration,
message='`concentration` must be non-negative'))
return assertions
def _sample_control_dependencies(self, x):
assertions = []
if tensorshape_util.is_fully_defined(x.shape[-2:]):
if not (tensorshape_util.dims(x.shape)[-2] ==
tensorshape_util.dims(x.shape)[-1] == self.dimension):
raise ValueError(
'Input dimension mismatch: expected [..., {}, {}], got {}'.
format(self.dimension, self.dimension,
tensorshape_util.dims(x.shape)))
elif self.validate_args:
msg = 'Input dimension mismatch: expected [..., {}, {}], got {}'.format(
self.dimension, self.dimension, tf.shape(x))
assertions.append(
assert_util.assert_equal(tf.shape(x)[-2],
self.dimension,
message=msg))
assertions.append(
assert_util.assert_equal(tf.shape(x)[-1],
self.dimension,
message=msg))
if self.validate_args and not self.input_output_cholesky:
assertions.append(
assert_util.assert_less_equal(
dtype_util.as_numpy_dtype(x.dtype)(-1),
x,
message='Correlations must be >= -1.',
summarize=30))
assertions.append(
assert_util.assert_less_equal(
x,
dtype_util.as_numpy_dtype(x.dtype)(1),
message='Correlations must be <= 1.',
summarize=30))
assertions.append(
assert_util.assert_near(
tf.linalg.diag_part(x),
dtype_util.as_numpy_dtype(x.dtype)(1),
message='Self-correlations must be = 1.',
summarize=30))
assertions.append(
assert_util.assert_near(
x,
tf.linalg.matrix_transpose(x),
message='Correlation matrices must be symmetric.',
summarize=30))
return assertions
def _assert_valid_sample(self, x):
if not self.validate_args:
return x
return distribution_util.with_dependencies([
assert_util.assert_non_positive(x),
assert_util.assert_near(
tf.zeros([], dtype=self.dtype), tf.reduce_logsumexp(x, axis=[-1])),
], x)
def _is_valid_correlation_cholesky(self, x):
if not self.validate_args:
return []
return [
assert_util.assert_near(
x,
tf.linalg.band_part(x, -1, 0),
message='Cholesky factors must be lower triangular.')
]
def _maybe_assert_valid_sample(self, x, dtype):
if not self.validate_args:
return x
one = tf.ones([], dtype=dtype)
return distribution_util.with_dependencies([
assert_util.assert_non_negative(x),
assert_util.assert_less_equal(x, one),
assert_util.assert_near(one, tf.reduce_sum(x, axis=[-1])),
], x)
def _maybe_assert_valid_sample(self, x):
"""Checks the validity of a sample."""
if not self.validate_args:
return x
return distribution_util.with_dependencies([
assert_util.assert_positive(x, message="samples must be positive"),
assert_util.assert_near(
tf.ones([], dtype=self.dtype),
tf.reduce_sum(input_tensor=x, axis=-1),
message="sample last-dimension must sum to `1`"),
], x)
def _is_valid_correlation_matrix(self, x):
if not self.validate_args or self.input_output_cholesky:
return []
return [
assert_util.assert_less_equal(
dtype_util.as_numpy_dtype(x.dtype)(-1),
x,
message='Correlations must be >= -1.'),
assert_util.assert_less_equal(
x,
dtype_util.as_numpy_dtype(x.dtype)(1),
message='Correlations must be <= 1.'),
assert_util.assert_near(tf.linalg.diag_part(x),
dtype_util.as_numpy_dtype(x.dtype)(1),
message='Self-correlations must be = 1.'),
assert_util.assert_near(
x,
tf.linalg.matrix_transpose(x),
message='Correlation matrices must be symmetric')
]
def _maybe_assert_valid_sample(self, x):
"""Checks the validity of a sample."""
if not self.validate_args:
return []
return [
assert_util.assert_positive(x, message='samples must be positive'),
assert_util.assert_near(
tf.ones([], dtype=self.dtype),
tf.reduce_sum(x, axis=-1),
message='sample last-dimension must sum to `1`'),
]
def _sample_control_dependencies(self, x):
"""Checks the validity of a sample."""
assertions = []
if not self.validate_args:
return assertions
assertions.append(assert_util.assert_non_negative(
x, message='Samples must be non-negative.'))
assertions.append(assert_util.assert_near(
tf.ones([], dtype=self.dtype),
tf.reduce_sum(x, axis=-1),
message='Sample last-dimension must sum to `1`.'))
return assertions
def _assert_ops(
self,
ode_fn,
initial_time,
initial_state,
solution_times,
previous_solver_state,
rtol,
atol,
first_step_size,
safety_factor,
min_step_size_factor,
max_step_size_factor,
max_num_steps,
solution_times_by_solver
):
"""Constructs dynamic assertions that validate input values to `_solve`."""
assert_ops = []
if self._validate_args is None:
return assert_ops
if solution_times_by_solver:
final_time = solution_times.final_time
assert_ops.append(
util.assert_positive(final_time - initial_time,
'final_time - initial_time'))
else:
assert_ops += [
util.assert_increasing(solution_times, 'solution_times'),
util.assert_nonnegative(solution_times[0] - initial_time,
'solution_times[0] - initial_time'),
]
if previous_solver_state is not None:
state_diff = initial_state - previous_solver_state.current_state
assert_states_match = assert_util.assert_near(
tf.norm(state_diff), 0., message='`previous_solver_state` does not '
'match the `initial_state`.')
assert_ops.append(assert_states_match)
if self._max_num_steps is not None:
assert_ops.append(util.assert_positive(max_num_steps, 'max_num_steps'))
assert_ops += [
util.assert_positive(rtol, 'rtol'),
util.assert_positive(atol, 'atol'),
util.assert_positive(first_step_size, 'first_step_size'),
util.assert_positive(safety_factor, 'safety_factor'),
util.assert_positive(
min_step_size_factor, 'min_step_size_factor'),
util.assert_positive(
max_step_size_factor, 'max_step_size_factor'),
]
derivative = ode_fn(initial_time, initial_state)
tf.nest.assert_same_structure(initial_state, derivative)
return assert_ops
def _validate_correlationness(self, x):
if not self.validate_args or self.input_output_cholesky:
return x
checks = [
assert_util.assert_less_equal(
dtype_util.as_numpy_dtype(x.dtype)(-1),
x,
message='Correlations must be >= -1.'),
assert_util.assert_less_equal(
x,
dtype_util.as_numpy_dtype(x.dtype)(1),
message='Correlations must be <= 1.'),
assert_util.assert_near(tf.linalg.diag_part(x),
dtype_util.as_numpy_dtype(x.dtype)(1),
message='Self-correlations must be = 1.'),
assert_util.assert_near(
x,
tf.linalg.matrix_transpose(x),
message='Correlation matrices must be symmetric')
]
with tf.control_dependencies(checks):
return tf.identity(x)
def maybe_assert_categorical_param_correctness(is_init, validate_args, probs,
logits):
"""Return assertions for `Categorical`-type distributions."""
assertions = []
# In init, we can always build shape and dtype checks because
# we assume shape doesn't change for Variable backed args.
if is_init:
x, name = (probs, 'probs') if logits is None else (logits, 'logits')
if not dtype_util.is_floating(x.dtype):
raise TypeError(
'Argument `{}` must having floating type.'.format(name))
msg = 'Argument `{}` must have rank at least 1.'.format(name)
ndims = tensorshape_util.rank(x.shape)
if ndims is not None:
if ndims < 1:
raise ValueError(msg)
elif validate_args:
x = tf.convert_to_tensor(x)
probs = x if logits is None else None # Retain tensor conversion.
logits = x if probs is None else None
assertions.append(
assert_util.assert_rank_at_least(x, 1, message=msg))
if not validate_args:
assert not assertions # Should never happen.
return []
if logits is not None:
if is_init != tensor_util.is_mutable(logits):
logits = tf.convert_to_tensor(logits)
assertions.extend(
distribution_util.assert_categorical_event_shape(logits))
if probs is not None:
if is_init != tensor_util.is_mutable(probs):
probs = tf.convert_to_tensor(probs)
assertions.extend([
assert_util.assert_non_negative(probs),
assert_util.assert_near(
tf.reduce_sum(probs, axis=-1),
np.array(1, dtype=dtype_util.as_numpy_dtype(probs.dtype)),
message='Argument `probs` must sum to 1.')
])
assertions.extend(
distribution_util.assert_categorical_event_shape(probs))
return assertions
def _sample_control_dependencies(self, x):
assertions = []
if not self.validate_args:
return assertions
assertions.append(assert_util.assert_non_positive(
x,
message=('Samples must be less than or equal to `0` for '
'`ExpRelaxedOneHotCategorical` or `1` for '
'`RelaxedOneHotCategorical`.')))
assertions.append(assert_util.assert_near(
tf.zeros([], dtype=self.dtype), tf.reduce_logsumexp(x, axis=[-1]),
message=('Final dimension of samples must sum to `0` for ''.'
'`ExpRelaxedOneHotCategorical` or `1` '
'for `RelaxedOneHotCategorical`.')))
return assertions
def _maybe_assert_valid_sample(self, samples):
"""Check counts for proper shape, values, then return tensor version."""
if not self.validate_args:
return samples
with tf.control_dependencies([
assert_util.assert_near(
1.,
tf.linalg.norm(samples, axis=-1),
message='samples must be unit length'),
assert_util.assert_equal(
tf.shape(samples)[-1:],
self.event_shape_tensor(),
message=('samples must have innermost dimension matching that of '
'`self.mean_direction`')),
]):
return tf.identity(samples)
def _sample_control_dependencies(self, samples):
inner_sample_dim = samples.shape[-1]
shape_msg = ('Samples must have innermost dimension matching that of '
'`self.dimension`. Found {}, expected {}'.format(
inner_sample_dim, self.dimension))
if inner_sample_dim is not None:
if self.dimension != inner_sample_dim:
raise ValueError(shape_msg)
assertions = []
if not self.validate_args:
return assertions
assertions.append(
assert_util.assert_near(tf.cast(1., dtype=self.dtype),
tf.linalg.norm(samples, axis=-1),
message='Samples must be unit length.'))
assertions.append(
assert_util.assert_equal(tf.shape(samples)[-1:],
self.dimension,
message=shape_msg))
return assertions
def _sample_control_dependencies(self, samples):
"""Check samples for proper shape and whether samples are unit vectors."""
inner_sample_dim = samples.shape[-1]
event_size = self.event_shape[-1]
shape_msg = ('Samples must have innermost dimension matching that of '
'`self.mean_direction`.')
if event_size is not None and inner_sample_dim is not None:
if event_size != inner_sample_dim:
raise ValueError(shape_msg)
assertions = []
if not self.validate_args:
return assertions
assertions.append(assert_util.assert_near(
1.,
tf.linalg.norm(samples, axis=-1),
message='Samples must be unit length.'))
assertions.append(assert_util.assert_equal(
tf.shape(samples)[-1:],
self.event_shape_tensor(),
message=shape_msg))
return assertions
def _parameter_control_dependencies(self, is_init):
"""Validate parameters."""
bw, bh, kd = None, None, None
try:
shape = tf.broadcast_static_shape(self.bin_widths.shape,
self.bin_heights.shape)
except ValueError as e:
raise ValueError('`bin_widths`, `bin_heights` must broadcast: {}'.format(
str(e)))
bin_sizes_shape = shape
try:
shape = tf.broadcast_static_shape(shape[:-1], self.knot_slopes.shape[:-1])
except ValueError as e:
raise ValueError(
'`bin_widths`, `bin_heights`, and `knot_slopes` must broadcast on '
'batch axes: {}'.format(str(e)))
assertions = []
if (tensorshape_util.is_fully_defined(bin_sizes_shape[-1:]) and
tensorshape_util.is_fully_defined(self.knot_slopes.shape[-1:])):
if tensorshape_util.rank(self.knot_slopes.shape) > 0:
num_interior_knots = tensorshape_util.dims(bin_sizes_shape)[-1] - 1
if tensorshape_util.dims(
self.knot_slopes.shape)[-1] not in (1, num_interior_knots):
raise ValueError(
'Innermost axis of non-scalar `knot_slopes` must broadcast with '
'{}; got {}.'.format(num_interior_knots, self.knot_slopes.shape))
elif self.validate_args:
if is_init != any(
tensor_util.is_ref(t)
for t in (self.bin_widths, self.bin_heights, self.knot_slopes)):
bw = tf.convert_to_tensor(self.bin_widths) if bw is None else bw
bh = tf.convert_to_tensor(self.bin_heights) if bh is None else bh
kd = _ensure_at_least_1d(self.knot_slopes) if kd is None else kd
shape = tf.broadcast_dynamic_shape(
tf.shape((bw + bh)[..., :-1]), tf.shape(kd))
assertions.append(
assert_util.assert_greater(
tf.shape(shape)[0],
tf.zeros([], dtype=shape.dtype),
message='`(bin_widths + bin_heights)[..., :-1]` must broadcast '
'with `knot_slopes` to at least 1-D.'))
if not self.validate_args:
assert not assertions
return assertions
if (is_init != tensor_util.is_ref(self.bin_widths) or
is_init != tensor_util.is_ref(self.bin_heights)):
bw = tf.convert_to_tensor(self.bin_widths) if bw is None else bw
bh = tf.convert_to_tensor(self.bin_heights) if bh is None else bh
assertions += [
assert_util.assert_near(
tf.reduce_sum(bw, axis=-1),
tf.reduce_sum(bh, axis=-1),
message='`sum(bin_widths, axis=-1)` must equal '
'`sum(bin_heights, axis=-1)`.'),
]
if is_init != tensor_util.is_ref(self.bin_widths):
bw = tf.convert_to_tensor(self.bin_widths) if bw is None else bw
assertions += [
assert_util.assert_positive(
bw, message='`bin_widths` must be positive.'),
]
if is_init != tensor_util.is_ref(self.bin_heights):
bh = tf.convert_to_tensor(self.bin_heights) if bh is None else bh
assertions += [
assert_util.assert_positive(
bh, message='`bin_heights` must be positive.'),
]
if is_init != tensor_util.is_ref(self.knot_slopes):
kd = _ensure_at_least_1d(self.knot_slopes) if kd is None else kd
assertions += [
assert_util.assert_positive(
kd, message='`knot_slopes` must be positive.'),
]
return assertions
def _sample_n(self, n, seed=None):
dim0_seed, otherdims_seed = samplers.split_seed(seed,
salt='von_mises_fisher')
# The sampling strategy relies on the fact that vMF variates are symmetric
# about the mean direction. Accordingly, if we have a sampling strategy for
# the away-from-mean angle, then we can uniformly sample the remaining
# dimensions on the S^{dim-2} sphere for , and rotate these samples from a
# (1, 0, 0, ..., 0)-mode distribution into the target orientation.
#
# This is easy to imagine on the 1-sphere (S^1; in 2-D space): sample a
# von-Mises distributed `x` value in [-1, 1], then uniformly select what
# amounts to a "up" or "down" additional degree of freedom after unit
# normalizing, followed by a final rotation to the desired mean direction
# from a basis of (1, 0).
#
# On S^2 (in 3-D), selecting a vMF `x` identifies a circle in `yz` on the
# unit sphere over which the distribution is uniform, in particular the
# circle where x = \hat{x} intersects the unit sphere. We pick a point on
# that circle, then rotate to the desired mean direction from a basis of
# (1, 0, 0).
mean_direction = tf.convert_to_tensor(self.mean_direction)
concentration = tf.convert_to_tensor(self.concentration)
event_dim = (
tf.compat.dimension_value(self.event_shape[0]) or
self._event_shape_tensor(mean_direction=mean_direction)[0])
sample_batch_shape = ps.concat([[n], self._batch_shape_tensor(
mean_direction=mean_direction, concentration=concentration)], axis=0)
dim = tf.cast(event_dim - 1, self.dtype)
if event_dim == 3:
samples_dim0 = self._sample_3d(n,
mean_direction=mean_direction,
concentration=concentration,
seed=dim0_seed)
else:
# Wood'94 provides a rejection algorithm to sample the x coordinate.
# Wood'94 definition of b:
# b = (-2 * kappa + tf.sqrt(4 * kappa**2 + dim**2)) / dim
# https://stats.stackexchange.com/questions/156729 suggests:
b = dim / (2 * concentration +
tf.sqrt(4 * concentration**2 + dim**2))
# TODO(bjp): Integrate any useful numerical tricks from hyperspherical VAE
# https://github.com/nicola-decao/s-vae-tf/
x = (1 - b) / (1 + b)
c = concentration * x + dim * tf.math.log1p(-x**2)
beta = beta_lib.Beta(dim / 2, dim / 2)
def cond_fn(w, should_continue, seed):
del w, seed
return tf.reduce_any(should_continue)
def body_fn(w, should_continue, seed):
"""While loop body for sampling the angle `w`."""
beta_seed, unif_seed, next_seed = samplers.split_seed(seed, n=3)
z = beta.sample(sample_shape=sample_batch_shape, seed=beta_seed)
# set_shape needed here because of b/139013403
tensorshape_util.set_shape(z, w.shape)
w = tf.where(should_continue,
(1. - (1. + b) * z) / (1. - (1. - b) * z),
w)
if not self.allow_nan_stats:
w = tf.debugging.check_numerics(w, 'w')
unif = samplers.uniform(
sample_batch_shape, seed=unif_seed, dtype=self.dtype)
# set_shape needed here because of b/139013403
tensorshape_util.set_shape(unif, w.shape)
should_continue = should_continue & (
concentration * w + dim * tf.math.log1p(-x * w) - c <
# Use log1p(-unif) to prevent log(0) and ensure that log(1) is
# possible.
tf.math.log1p(-unif))
return w, should_continue, next_seed
w = tf.zeros(sample_batch_shape, dtype=self.dtype)
should_continue = tf.ones(sample_batch_shape, dtype=tf.bool)
samples_dim0, _, _ = tf.while_loop(
cond=cond_fn, body=body_fn,
loop_vars=(w, should_continue, dim0_seed))
samples_dim0 = samples_dim0[..., tf.newaxis]
if not self._allow_nan_stats:
# Verify samples are w/in -1, 1, with useful error output tensors (top
# value rather than all values).
with tf.control_dependencies([
assert_util.assert_less_equal(
samples_dim0,
dtype_util.as_numpy_dtype(self.dtype)(1.01)),
assert_util.assert_greater_equal(
samples_dim0,
dtype_util.as_numpy_dtype(self.dtype)(-1.01)),
]):
samples_dim0 = tf.identity(samples_dim0)
samples_otherdims_shape = ps.concat([sample_batch_shape, [event_dim - 1]],
axis=0)
unit_otherdims = tf.math.l2_normalize(
samplers.normal(
samples_otherdims_shape, seed=otherdims_seed, dtype=self.dtype),
axis=-1)
samples = tf.concat([
samples_dim0, # we must avoid sqrt(1 - (>1)**2)
tf.sqrt(tf.maximum(1 - samples_dim0**2, 0.)) * unit_otherdims
], axis=-1)
samples = tf.math.l2_normalize(samples, axis=-1)
if not self.allow_nan_stats:
samples = tf.debugging.check_numerics(samples, 'samples')
# Runtime assert that samples are unit length.
if not self.allow_nan_stats:
worst, _ = tf.math.top_k(
tf.reshape(tf.abs(1 - tf.linalg.norm(samples, axis=-1)), [-1]))
with tf.control_dependencies([
assert_util.assert_near(
dtype_util.as_numpy_dtype(self.dtype)(0),
worst,
atol=1e-4,
summarize=100)
]):
samples = tf.identity(samples)
# The samples generated are symmetric around a mode at (1, 0, 0, ...., 0).
# Now, we move the mode to `self.mean_direction` using a rotation matrix.
if not self.allow_nan_stats:
# Assert that the basis vector rotates to the mean direction, as expected.
basis = tf.cast(tf.concat([[1.], tf.zeros([event_dim - 1])], axis=0),
self.dtype)
with tf.control_dependencies([
assert_util.assert_less(
tf.linalg.norm(
self._rotate(basis, mean_direction=mean_direction) -
mean_direction, axis=-1),
dtype_util.as_numpy_dtype(self.dtype)(1e-5))
]):
return self._rotate(samples, mean_direction=mean_direction)
return self._rotate(samples, mean_direction=mean_direction)
def __init__(self,
mean_direction,
concentration,
validate_args=False,
allow_nan_stats=True,
name='VonMisesFisher'):
"""Creates a new `VonMisesFisher` instance.
Args:
mean_direction: Floating-point `Tensor` with shape [B1, ... Bn, D].
A unit vector indicating the mode of the distribution, or the
unit-normalized direction of the mean. (This is *not* in general the
mean of the distribution; the mean is not generally in the support of
the distribution.) NOTE: `D` is currently restricted to <= 5.
concentration: Floating-point `Tensor` having batch shape [B1, ... Bn]
broadcastable with `mean_direction`. The level of concentration of
samples around the `mean_direction`. `concentration=0` indicates a
uniform distribution over the unit hypersphere, and `concentration=+inf`
indicates a `Deterministic` distribution (delta function) at
`mean_direction`.
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.
Raises:
ValueError: For known-bad arguments, i.e. unsupported event dimension.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([mean_direction, concentration],
tf.float32)
mean_direction = tf.convert_to_tensor(
mean_direction, name='mean_direction', dtype=dtype)
concentration = tf.convert_to_tensor(
concentration, name='concentration', dtype=dtype)
assertions = [
assert_util.assert_non_negative(
concentration, message='`concentration` must be non-negative'),
assert_util.assert_greater(
tf.shape(mean_direction)[-1],
1,
message='`mean_direction` may not have scalar event shape'),
assert_util.assert_near(
1.,
tf.linalg.norm(mean_direction, axis=-1),
message='`mean_direction` must be unit-length')
] if validate_args else []
static_event_dim = tf.compat.dimension_value(
tensorshape_util.with_rank_at_least(mean_direction.shape, 1)[-1])
if static_event_dim is not None and static_event_dim > 5:
raise ValueError('vMF ndims > 5 is not currently supported')
elif validate_args:
assertions += [
assert_util.assert_less_equal(
tf.shape(mean_direction)[-1],
5,
message='vMF ndims > 5 is not currently supported')
]
with tf.control_dependencies(assertions):
self._mean_direction = tf.identity(mean_direction)
self._concentration = tf.identity(concentration)
dtype_util.assert_same_float_dtype(
[self._mean_direction, self._concentration])
# mean_direction is always reparameterized.
# concentration is only for event_dim==3, via an inversion sampler.
reparameterization_type = (
reparameterization.FULLY_REPARAMETERIZED
if static_event_dim == 3 else
reparameterization.NOT_REPARAMETERIZED)
super(VonMisesFisher, self).__init__(
dtype=self._concentration.dtype,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
reparameterization_type=reparameterization_type,
parameters=parameters,
name=name)
def __init__(self,
loc=None,
covariance_matrix=None,
validate_args=False,
allow_nan_stats=True,
name='MultivariateNormalFullCovariance'):
"""Construct Multivariate Normal distribution on `R^k`.
The `batch_shape` is the broadcast shape between `loc` and
`covariance_matrix` arguments.
The `event_shape` is given by last dimension of the matrix implied by
`covariance_matrix`. The last dimension of `loc` (if provided) must
broadcast with this.
A non-batch `covariance_matrix` matrix is a `k x k` symmetric positive
definite matrix. In other words it is (real) symmetric with all eigenvalues
strictly positive.
Additional leading dimensions (if any) will index batches.
Args:
loc: Floating-point `Tensor`. If this is set to `None`, `loc` is
implicitly `0`. When specified, may have shape `[B1, ..., Bb, k]` where
`b >= 0` and `k` is the event size.
covariance_matrix: Floating-point, symmetric positive definite `Tensor` of
same `dtype` as `loc`. The strict upper triangle of `covariance_matrix`
is ignored, so if `covariance_matrix` is not symmetric no error will be
raised (unless `validate_args is True`). `covariance_matrix` has shape
`[B1, ..., Bb, k, k]` where `b >= 0` and `k` is the event size.
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.
Raises:
ValueError: if neither `loc` nor `covariance_matrix` are specified.
"""
parameters = dict(locals())
# Convert the covariance_matrix up to a scale_tril and call MVNTriL.
with tf.name_scope(name) as name:
with tf.name_scope('init'):
dtype = dtype_util.common_dtype([loc, covariance_matrix],
tf.float32)
loc = loc if loc is None else tf.convert_to_tensor(
loc, name='loc', dtype=dtype)
if covariance_matrix is None:
scale_tril = None
else:
covariance_matrix = tf.convert_to_tensor(
covariance_matrix,
name='covariance_matrix',
dtype=dtype)
if validate_args:
covariance_matrix = distribution_util.with_dependencies(
[
assert_util.assert_near(
covariance_matrix,
tf.linalg.matrix_transpose(
covariance_matrix),
message='Matrix was not symmetric')
], covariance_matrix)
# No need to validate that covariance_matrix is non-singular.
# LinearOperatorLowerTriangular has an assert_non_singular method that
# is called by the Bijector.
# However, cholesky() ignores the upper triangular part, so we do need
# to separately assert symmetric.
scale_tril = tf.linalg.cholesky(covariance_matrix)
super(MultivariateNormalFullCovariance,
self).__init__(loc=loc,
scale_tril=scale_tril,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
name=name)
self._parameters = parameters