def _parameter_control_dependencies(self, is_init):
if not self.validate_args:
return []
assertions = []
peak = None
if is_init != tensor_util.is_ref(self.low):
peak = tf.convert_to_tensor(self.peak)
assertions.append(
assert_util.assert_greater(
peak,
self.low,
message='`peak` must be greater than `low`.'))
if is_init != tensor_util.is_ref(self.high):
peak = tf.convert_to_tensor(self.peak) if peak is None else peak
assertions.append(
assert_util.assert_greater(
self.high,
peak,
message='`high` must be greater than `peak`.'))
if is_init != tensor_util.is_ref(self.temperature):
assertions.append(
assert_util.assert_positive(
self.temperature,
message='`temperature` must be positive.'))
return assertions
def _maybe_check_valid_map_values(map_values, validate_args):
"""Validate `map_values` if `validate_args`==True."""
assertions = []
message = 'Rank of map_values must be 1.'
if tensorshape_util.rank(map_values.shape) is not None:
if tensorshape_util.rank(map_values.shape) != 1:
raise ValueError(message)
elif validate_args:
assertions.append(
assert_util.assert_rank(map_values, 1, message=message))
message = 'Size of map_values must be greater than 0.'
if tensorshape_util.num_elements(map_values.shape) is not None:
if tensorshape_util.num_elements(map_values.shape) == 0:
raise ValueError(message)
elif validate_args:
assertions.append(
assert_util.assert_greater(tf.size(map_values), 0,
message=message))
if validate_args:
assertions.append(
assert_util.assert_equal(
tf.math.is_strictly_increasing(map_values),
True,
message='map_values is not strictly increasing.'))
return assertions
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 _parameter_control_dependencies(self, is_init):
if not self.validate_args:
return []
assertions = []
low = None
high = None
if is_init != tensor_util.is_ref(self.low):
low = tf.convert_to_tensor(self.low)
assertions.append(
assert_util.assert_finite(low, message='`low` is not finite'))
if is_init != tensor_util.is_ref(self.high):
high = tf.convert_to_tensor(self.high)
assertions.append(
assert_util.assert_finite(high, message='`high` is not finite'))
if is_init != tensor_util.is_ref(self.loc):
assertions.append(
assert_util.assert_finite(self.loc, message='`loc` is not finite'))
if is_init != tensor_util.is_ref(self.scale):
scale = tf.convert_to_tensor(self.scale)
assertions.extend([
assert_util.assert_positive(
scale, message='`scale` must be positive'),
assert_util.assert_finite(scale, message='`scale` is not finite'),
])
if (is_init != tensor_util.is_ref(self.low) or
is_init != tensor_util.is_ref(self.high)):
low = tf.convert_to_tensor(self.low) if low is None else low
high = tf.convert_to_tensor(self.high) if high is None else high
assertions.append(
assert_util.assert_greater(
high,
low,
message='TruncatedCauchy not defined when `low >= high`.'))
return assertions
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 _maybe_validate_split_sizes(self):
"""Validations for `split_sizes` property."""
assertions = []
split_sizes = tf.convert_to_tensor(self.split_sizes)
split_sizes_ = tf.get_static_value(split_sizes)
# Ensure `split_sizes` has no more than one unknown split size (=-1).
message = '`{}` elements must have at most one `-1`.'
if split_sizes_ is not None:
if sum(split_sizes_ == -1) > 1:
raise ValueError(message.format(split_sizes))
elif self.validate_args:
assertions.append(
assert_util.assert_less(tf.reduce_sum(
tf.cast(tf.equal(split_sizes, -1), tf.int32)),
2,
message=message.format(split_sizes)))
message = '`{}` elements must be either non-negative integers or `-1`.'
if split_sizes_ is not None:
if np.any(split_sizes_ < -1):
raise ValueError(message.format(split_sizes))
elif self.validate_args:
assertions.append(
assert_util.assert_greater(
split_sizes, -2, message=message.format(split_sizes)))
return assertions
def _inverse_event_shape_tensor(self, output_shape):
if self.validate_args:
# It is not possible for a negative shape so we need only check <= 1.
is_greater_one = assert_util.assert_greater(
output_shape[-1], 1, message="Need last dimension greater than 1.")
output_shape = distribution_util.with_dependencies(
[is_greater_one], output_shape)
return tf.concat([output_shape[:-1], [output_shape[-1] - 1]], axis=0)
def _parameter_control_dependencies(self, is_init):
if not self.validate_args:
return []
assertions = []
if is_init != tensor_util.is_ref(self.power):
assertions.append(assert_util.assert_greater(
self.power, np.ones([], dtype_util.as_numpy_dtype(self.power.dtype)),
message='`power` must be greater than 1.'))
return assertions
def _inverse_event_shape_tensor(self, output_shape):
if self.validate_args:
# It is not possible for a negative shape so we need only check <= 1.
dependencies = [assert_util.assert_greater(
output_shape[-1], 1, message="Need last dimension greater than 1.")]
else:
dependencies = []
with tf.control_dependencies(dependencies):
return tf.concat([output_shape[:-1], [output_shape[-1] - 1]], axis=0)
def _log_prob(self, x):
with tf.control_dependencies([
assert_util.assert_greater(
x, tf.cast(0., x.dtype.base_dtype), message="x must be positive.")
] if self.validate_args else []):
return (0.5 * (tf.math.log(self.concentration) - np.log(2. * np.pi) -
3. * tf.math.log(x)) + (-self.concentration *
(x - self.loc)**2.) /
(2. * self.loc**2. * x))
def _assertions(self, t):
if not self.validate_args:
return []
return [
assert_util.assert_greater(
t[..., 1:],
t[..., :-1],
message=
'Inverse transformation input must be strictly increasing.')
]
def _parameter_control_dependencies(self, is_init):
if not self.validate_args:
return []
assertions = []
if is_init != tensor_util.is_ref(self.df):
assertions.append(
assert_util.assert_greater(
self.df, dtype_util.as_numpy_dtype(self.df.dtype)(2.),
message='`df` must be greater than 2.'))
return assertions
def _parameter_control_dependencies(self, is_init):
if not self.validate_args or self.low is None or self.high is None:
return []
assertions = []
if is_init != (tensor_util.is_ref(self.low) or
tensor_util.is_ref(self.high)):
assertions.append(assert_util.assert_greater(
self.high, self.low,
message='Argument `high` must be greater than `low`.'))
return assertions
def _assert_valid_inverse_input(self, y):
assertions = []
if self.validate_args and self.low is not None:
assertions += [assert_util.assert_greater(
y, self.low,
message='Input must be greater than `low`.')]
if self.validate_args and self.high is not None:
assertions += [assert_util.assert_less(
y, self.high,
message='Input must be less than `high`.')]
return assertions
def _cdf(self, x):
with tf.control_dependencies([
assert_util.assert_greater(x,
tf.cast(0., x.dtype.base_dtype),
message="x must be positive.")
] if self.validate_args else []):
return (special_math.ndtr(
((self.concentration / x)**0.5 * (x / self.loc - 1.))) +
tf.exp(2. * self.concentration / self.loc) *
special_math.ndtr(-(self.concentration / x)**0.5 *
(x / self.loc + 1)))
def _inverse_event_shape_tensor(self, output_shape):
input_shape = self._forward_event_shape_tensor(
output_shape, is_inverse=True)
if not self.validate_args:
return input_shape
assertions = [
assert_util.assert_greater(
input_shape, -1,
message='Invalid inverse shape; found negative size.')
]
with tf.control_dependencies(assertions):
return tf.identity(input_shape)
def _assertions(self, t):
if not self.validate_args:
return []
return [
assert_util.assert_greater(
t,
dtype_util.as_numpy_dtype(t.dtype)(-1),
message="Inverse transformation input must be greater than -1."
),
assert_util.assert_less(
t,
dtype_util.as_numpy_dtype(t.dtype)(1),
message="Inverse transformation input must be less than 1.")
]
def _parameter_control_dependencies(self, is_init):
if is_init:
# Check that locs and slopes have the same last dimension.
if (self.locs.shape is not None and self.locs.shape[-1] is not None
and self.slopes.shape is not None
and self.slopes.shape[-1] is not None):
k = self.locs.shape[-1]
l = self.slopes.shape[-1]
if not (k == l or k == 1 or l == 1):
raise ValueError(
'Expect that `locs` and `slopes` are broadcastable.')
if self.locs.shape is not None and self.locs.shape[-1] is not None:
if not (len(self.kernels) == self.locs.shape[-1] + 1
or self.locs.shape[-1] == 1):
raise ValueError(
'Expect that `locs` has last dimension `1` or `N - 1` where '
f'`N` is the number of kernels, but got {self.locs.shape[-1]}'
)
if self.slopes.shape is not None and self.slopes.shape[
-1] is not None:
if not (len(self.kernels) == self.slopes.shape[-1] + 1
or self.slopes.shape[-1] == 1):
raise ValueError(
'Expect that `slopes` has last dimension `1` or `N - 1` where '
f'`N` is the number of kernels, but got {self.slopes.shape[-1]}'
)
assertions = []
if not self.validate_args:
return assertions
if is_init != tensor_util.is_ref(self.locs):
locs = tf.convert_to_tensor(self.locs)
assertions.append(
assert_util.assert_greater(
locs[..., 1:],
locs[..., :-1],
message='Expect that elements of `locs` are ascending.'))
if is_init != tensor_util.is_ref(self.slopes):
slopes = tf.convert_to_tensor(self.slopes)
assertions.append(
assert_util.assert_positive(
slopes, message='`slopes` must be positive.'))
return assertions
def _maybe_validate_target_accept_prob(target_accept_prob, validate_args):
"""Validates that target_accept_prob is in (0, 1)."""
if not validate_args:
return target_accept_prob
assertions = [
assert_util.assert_greater(
target_accept_prob,
tf.zeros([], dtype=target_accept_prob.dtype),
message='`target_accept_prob` must be > 0.'),
assert_util.assert_less(target_accept_prob,
tf.ones([], dtype=target_accept_prob.dtype),
message='`target_accept_prob` must be < 1.')
]
with tf.control_dependencies(assertions):
return tf.identity(target_accept_prob)
def _maybe_assert_valid_y(self, y):
if not self.validate_args:
return y
is_valid = [
assert_util.assert_greater(
y,
dtype_util.as_numpy_dtype(y.dtype)(-1),
message="Inverse transformation input must be greater than -1."
),
assert_util.assert_less(
y,
dtype_util.as_numpy_dtype(y.dtype)(1),
message="Inverse transformation input must be less than 1.")
]
return distribution_util.with_dependencies(is_valid, y)
def _parameter_control_dependencies(self, is_init):
assertions = []
sample_shape = None # Memoize concretization.
# Check valid shape.
ndims_ = tensorshape_util.rank(self.sample_shape.shape)
if is_init != (ndims_ is None):
msg = 'Argument `sample_shape` must be either a scalar or a vector.'
if ndims_ is not None:
if ndims_ > 1:
raise ValueError(msg)
elif self.validate_args:
if sample_shape is None:
sample_shape = tf.convert_to_tensor(self.sample_shape)
assertions.append(
assert_util.assert_less(tf.rank(sample_shape),
2,
message=msg))
# Check valid dtype.
if is_init: # No xor check because `dtype` cannot change.
dtype_ = self.sample_shape.dtype
if dtype_ is None:
if sample_shape is None:
sample_shape = tf.convert_to_tensor(self.sample_shape)
dtype_ = sample_shape.dtype
if dtype_util.base_dtype(dtype_) not in {tf.int32, tf.int64}:
raise TypeError(
'Argument `sample_shape` must be integer type; '
'saw {}.'.format(dtype_util.name(dtype_)))
# Check valid "value".
if is_init != tensor_util.is_ref(self.sample_shape):
sample_shape_ = tf.get_static_value(self.sample_shape)
msg = 'Argument `sample_shape` must have non-negative values.'
if sample_shape_ is not None:
if np.any(np.array(sample_shape_) < 0):
raise ValueError('{} Saw: {}'.format(msg, sample_shape_))
elif self.validate_args:
if sample_shape is None:
sample_shape = tf.convert_to_tensor(self.sample_shape)
assertions.append(
assert_util.assert_greater(sample_shape, -1, message=msg))
return assertions
def _maybe_check_valid_shape(shape, validate_args):
"""Check that a shape Tensor is int-type and otherwise sane."""
if not dtype_util.is_integer(shape.dtype):
raise TypeError('`{}` dtype (`{}`) should be `int`-like.'.format(
shape, dtype_util.name(shape.dtype)))
assertions = []
message = '`{}` rank should be <= 1.'
if tensorshape_util.rank(shape.shape) is not None:
if tensorshape_util.rank(shape.shape) > 1:
raise ValueError(message.format(shape))
elif validate_args:
assertions.append(
assert_util.assert_less(tf.rank(shape),
2,
message=message.format(shape)))
shape_ = tf.get_static_value(shape)
message = '`{}` elements must have at most one `-1`.'
if shape_ is not None:
if sum(shape_ == -1) > 1:
raise ValueError(message.format(shape))
elif validate_args:
assertions.append(
assert_util.assert_less(tf.reduce_sum(
tf.cast(tf.equal(shape, -1), tf.int32)),
2,
message=message.format(shape)))
message = '`{}` elements must be either positive integers or `-1`.'
if shape_ is not None:
if np.any(shape_ < -1):
raise ValueError(message.format(shape))
elif validate_args:
assertions.append(
assert_util.assert_greater(shape,
-2,
message=message.format(shape)))
return assertions
def _parameter_control_dependencies(self, is_init):
if not self.validate_args:
return []
assertions = []
if is_init != tensor_util.is_ref(self.mass):
assertions.append(
assert_util.assert_positive(
self.mass, message='Argument `mass` must be positive.'))
if is_init != tensor_util.is_ref(self.width):
assertions.append(
assert_util.assert_positive(
self.width, message='Argument `width` must be positive.'))
if is_init != tensor_util.is_ref(self.smin):
assertions.append(
assert_util.assert_non_negative(
self.smin,
message='Argument `smin` must be positive or zero.'))
if is_init != tensor_util.is_ref(self.smax):
assertions.append(
assert_util.assert_greater(
self.smax,
self.smin,
message='Argument `smax` must be larger than `smin`.'))
return assertions
def __init__(self,
initial_distribution,
transition_distribution,
observation_distribution,
num_steps,
validate_args=False,
allow_nan_stats=True,
name="HiddenMarkovModel"):
"""Initialize hidden Markov model.
Args:
initial_distribution: A `Categorical`-like instance.
Determines probability of first hidden state in Markov chain.
The number of categories must match the number of categories of
`transition_distribution` as well as both the rightmost batch
dimension of `transition_distribution` and the rightmost batch
dimension of `observation_distribution`.
transition_distribution: A `Categorical`-like instance.
The rightmost batch dimension indexes the probability distribution
of each hidden state conditioned on the previous hidden state.
observation_distribution: A `tfp.distributions.Distribution`-like
instance. The rightmost batch dimension indexes the distribution
of each observation conditioned on the corresponding hidden state.
num_steps: The number of steps taken in Markov chain. A python `int`.
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: "HiddenMarkovModel".
Raises:
ValueError: if `num_steps` is not at least 1.
ValueError: if `initial_distribution` does not have scalar `event_shape`.
ValueError: if `transition_distribution` does not have scalar
`event_shape.`
ValueError: if `transition_distribution` and `observation_distribution`
are fully defined but don't have matching rightmost dimension.
"""
parameters = dict(locals())
# pylint: disable=protected-access
with tf.name_scope(name) as name:
self._runtime_assertions = [] # pylint: enable=protected-access
num_steps = tf.convert_to_tensor(value=num_steps, name="num_steps")
if validate_args:
self._runtime_assertions += [
assert_util.assert_equal(
tf.rank(num_steps),
0,
message="`num_steps` must be a scalar")
]
self._runtime_assertions += [
assert_util.assert_greater_equal(
num_steps,
1,
message="`num_steps` must be at least 1.")
]
self._initial_distribution = initial_distribution
self._observation_distribution = observation_distribution
self._transition_distribution = transition_distribution
if (initial_distribution.event_shape is not None
and tensorshape_util.rank(
initial_distribution.event_shape) != 0):
raise ValueError(
"`initial_distribution` must have scalar `event_dim`s")
elif validate_args:
self._runtime_assertions += [
assert_util.assert_equal(
tf.shape(initial_distribution.event_shape_tensor())[0],
0,
message="`initial_distribution` must have scalar"
"`event_dim`s")
]
if (transition_distribution.event_shape is not None
and tensorshape_util.rank(
transition_distribution.event_shape) != 0):
raise ValueError(
"`transition_distribution` must have scalar `event_dim`s")
elif validate_args:
self._runtime_assertions += [
assert_util.assert_equal(
tf.shape(
transition_distribution.event_shape_tensor())[0],
0,
message="`transition_distribution` must have scalar"
"`event_dim`s")
]
if (tensorshape_util.dims(transition_distribution.batch_shape)
is not None and tensorshape_util.rank(
transition_distribution.batch_shape) == 0):
raise ValueError(
"`transition_distribution` can't have scalar batches")
elif validate_args:
self._runtime_assertions += [
assert_util.assert_greater(
tf.size(transition_distribution.batch_shape_tensor()),
0,
message="`transition_distribution` can't have scalar "
"batches")
]
if (tensorshape_util.dims(observation_distribution.batch_shape)
is not None and tensorshape_util.rank(
observation_distribution.batch_shape) == 0):
raise ValueError(
"`observation_distribution` can't have scalar batches")
elif validate_args:
self._runtime_assertions += [
assert_util.assert_greater(
tf.size(observation_distribution.batch_shape_tensor()),
0,
message="`observation_distribution` can't have scalar "
"batches")
]
# Infer number of hidden states and check consistency
# between transitions and observations
with tf.control_dependencies(self._runtime_assertions):
self._num_states = (
(tensorshape_util.dims(transition_distribution.batch_shape)
is not None and tensorshape_util.as_list(
transition_distribution.batch_shape)[-1])
or transition_distribution.batch_shape_tensor()[-1])
observation_states = (
(tensorshape_util.dims(
observation_distribution.batch_shape) is not None
and tensorshape_util.as_list(
observation_distribution.batch_shape)[-1])
or observation_distribution.batch_shape_tensor()[-1])
if (tf.is_tensor(self._num_states)
or tf.is_tensor(observation_states)):
if validate_args:
self._runtime_assertions += [
assert_util.assert_equal(
self._num_states,
observation_states,
message="`transition_distribution` and "
"`observation_distribution` must agree on "
"last dimension of batch size")
]
elif self._num_states != observation_states:
raise ValueError("`transition_distribution` and "
"`observation_distribution` must agree on "
"last dimension of batch size")
self._log_init = _extract_log_probs(self._num_states,
initial_distribution)
self._log_trans = _extract_log_probs(self._num_states,
transition_distribution)
self._num_steps = num_steps
self._num_states = tf.shape(self._log_init)[-1]
self._underlying_event_rank = tf.size(
self._observation_distribution.event_shape_tensor())
num_steps_ = tf.get_static_value(num_steps)
if num_steps_ is not None:
self.static_event_shape = tf.TensorShape([
num_steps_
]).concatenate(self._observation_distribution.event_shape)
else:
self.static_event_shape = None
with tf.control_dependencies(self._runtime_assertions):
self.static_batch_shape = tf.broadcast_static_shape(
self._initial_distribution.batch_shape,
tf.broadcast_static_shape(
self._transition_distribution.batch_shape[:-1],
self._observation_distribution.batch_shape[:-1]))
# pylint: disable=protected-access
super(HiddenMarkovModel, self).__init__(
dtype=self._observation_distribution.dtype,
reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
name=name)
# pylint: enable=protected-access
self._parameters = parameters
def _potential_scale_reduction_single_state(state, independent_chain_ndims,
split_chains, validate_args):
"""potential_scale_reduction for one single state `Tensor`."""
with tf.name_scope('potential_scale_reduction_single_state'):
# We assume exactly one leading dimension indexes e.g. correlated samples
# from each Markov chain.
state = tf.convert_to_tensor(state, name='state')
n_samples_ = tf.compat.dimension_value(state.shape[0])
if n_samples_ is not None: # If available statically.
if split_chains and n_samples_ < 4:
raise ValueError(
'Must provide at least 4 samples when splitting chains. '
'Found {}'.format(n_samples_))
if not split_chains and n_samples_ < 2:
raise ValueError(
'Must provide at least 2 samples. Found {}'.format(
n_samples_))
elif validate_args:
if split_chains:
assertions = [
assert_util.assert_greater(
tf.shape(state)[0],
4,
message=
'Must provide at least 4 samples when splitting chains.'
)
]
with tf.control_dependencies(assertions):
state = tf.identity(state)
else:
assertions = [
assert_util.assert_greater(
tf.shape(state)[0],
2,
message='Must provide at least 2 samples.')
]
with tf.control_dependencies(assertions):
state = tf.identity(state)
# Define so it's not a magic number.
# Warning! `if split_chains` logic assumes this is 1!
sample_ndims = 1
if split_chains:
# Split the sample dimension in half, doubling the number of
# independent chains.
# For odd number of samples, keep all but the last sample.
state_shape = prefer_static.shape(state)
n_samples = state_shape[0]
state = state[:n_samples - n_samples % 2]
# Suppose state = [0, 1, 2, 3, 4, 5]
# Step 1: reshape into [[0, 1, 2], [3, 4, 5]]
# E.g. reshape states of shape [a, b] into [2, a//2, b].
state = tf.reshape(
state,
prefer_static.concat([[2, n_samples // 2], state_shape[1:]],
axis=0))
# Step 2: Put the size `2` dimension in the right place to be treated as a
# chain, changing [[0, 1, 2], [3, 4, 5]] into [[0, 3], [1, 4], [2, 5]],
# reshaping [2, a//2, b] into [a//2, 2, b].
state = tf.transpose(
a=state,
perm=prefer_static.concat(
[[1, 0], tf.range(2, tf.rank(state))], axis=0))
# We're treating the new dim as indexing 2 chains, so increment.
independent_chain_ndims += 1
sample_axis = tf.range(0, sample_ndims)
chain_axis = tf.range(sample_ndims,
sample_ndims + independent_chain_ndims)
sample_and_chain_axis = tf.range(
0, sample_ndims + independent_chain_ndims)
n = _axis_size(state, sample_axis)
m = _axis_size(state, chain_axis)
# In the language of Brooks and Gelman (1998),
# B / n is the between chain variance, the variance of the chain means.
# W is the within sequence variance, the mean of the chain variances.
b_div_n = _reduce_variance(tf.reduce_mean(state,
axis=sample_axis,
keepdims=True),
sample_and_chain_axis,
biased=False)
w = tf.reduce_mean(_reduce_variance(state,
sample_axis,
keepdims=True,
biased=True),
axis=sample_and_chain_axis)
# sigma^2_+ is an estimate of the true variance, which would be unbiased if
# each chain was drawn from the target. c.f. "law of total variance."
sigma_2_plus = w + b_div_n
return ((m + 1.) / m) * sigma_2_plus / w - (n - 1.) / (m * n)
def _effective_sample_size_single_state(states, filter_beyond_lag,
filter_threshold,
filter_beyond_positive_pairs,
cross_chain_dims, validate_args):
"""ESS computation for one single Tensor argument."""
with tf.name_scope('effective_sample_size_single_state'):
states = tf.convert_to_tensor(states, name='states')
dt = states.dtype
# filter_beyond_lag == None ==> auto_corr is the full sequence.
auto_cov = stats.auto_correlation(states,
axis=0,
max_lags=filter_beyond_lag,
normalize=False)
n = _axis_size(states, axis=0)
if cross_chain_dims is not None:
num_chains = _axis_size(states, cross_chain_dims)
num_chains_ = tf.get_static_value(num_chains)
assertions = []
msg = (
'When `cross_chain_dims` is not `None`, there must be > 1 chain '
'in `states`.')
if num_chains_ is not None:
if num_chains_ < 2:
raise ValueError(msg)
elif validate_args:
assertions.append(
assert_util.assert_greater(num_chains, 1., message=msg))
with tf.control_dependencies(assertions):
# We're computing the R[k] from equation 10 of Vehtari et al.
# (2019):
#
# R[k] := 1 - (W - 1/C * Sum_{c=1}^C s_c**2 R[k, c]) / (var^+),
#
# where:
# C := number of chains
# N := length of chains
# x_hat[c] := 1 / N Sum_{n=1}^N x[n, c], chain mean.
# x_hat := 1 / C Sum_{c=1}^C x_hat[c], overall mean.
# W := 1/C Sum_{c=1}^C s_c**2, within-chain variance.
# B := N / (C - 1) Sum_{c=1}^C (x_hat[c] - x_hat)**2, between chain
# variance.
# s_c**2 := 1 / (N - 1) Sum_{n=1}^N (x[n, c] - x_hat[c])**2, chain
# variance
# R[k, m] := auto_corr[k, m, ...], auto-correlation indexed by chain.
# var^+ := (N - 1) / N * W + B / N
cross_chain_dims = prefer_static.non_negative_axis(
cross_chain_dims, prefer_static.rank(states))
# B / N
between_chain_variance_div_n = _reduce_variance(
tf.reduce_mean(states, axis=0),
biased=False, # This makes the denominator be C - 1.
axis=cross_chain_dims - 1)
# W * (N - 1) / N
biased_within_chain_variance = tf.reduce_mean(
auto_cov[0], cross_chain_dims - 1)
# var^+
approx_variance = (biased_within_chain_variance +
between_chain_variance_div_n)
# 1/C * Sum_{c=1}^C s_c**2 R[k, c]
mean_auto_cov = tf.reduce_mean(auto_cov, cross_chain_dims)
auto_corr = 1. - (biased_within_chain_variance -
mean_auto_cov) / approx_variance
else:
auto_corr = auto_cov / auto_cov[:1]
num_chains = 1
# With R[k] := auto_corr[k, ...],
# ESS = N / {1 + 2 * Sum_{k=1}^N R[k] * (N - k) / N}
# = N / {-1 + 2 * Sum_{k=0}^N R[k] * (N - k) / N} (since R[0] = 1)
# approx N / {-1 + 2 * Sum_{k=0}^M R[k] * (N - k) / N}
# where M is the filter_beyond_lag truncation point chosen above.
# Get the factor (N - k) / N, and give it shape [M, 1,...,1], having total
# ndims the same as auto_corr
k = tf.range(0., _axis_size(auto_corr, axis=0))
nk_factor = (n - k) / n
if tensorshape_util.rank(auto_corr.shape) is not None:
new_shape = [-1
] + [1] * (tensorshape_util.rank(auto_corr.shape) - 1)
else:
new_shape = tf.concat(
([-1], tf.ones([tf.rank(auto_corr) - 1], dtype=tf.int32)),
axis=0)
nk_factor = tf.reshape(nk_factor, new_shape)
weighted_auto_corr = nk_factor * auto_corr
if filter_beyond_positive_pairs:
def _sum_pairs(x):
x_len = tf.shape(x)[0]
# For odd sequences, we drop the final value.
x = x[:x_len - x_len % 2]
new_shape = tf.concat(
[[x_len // 2, 2], tf.shape(x)[1:]], axis=0)
return tf.reduce_sum(tf.reshape(x, new_shape), 1)
# Pairwise sums are all positive for auto-correlation spectra derived from
# reversible MCMC chains.
# E.g. imagine the pairwise sums are [0.2, 0.1, -0.1, -0.2]
# Step 1: mask = [False, False, True, True]
mask = _sum_pairs(auto_corr) < 0.
# Step 2: mask = [0, 0, 1, 1]
mask = tf.cast(mask, dt)
# Step 3: mask = [0, 0, 1, 2]
mask = tf.cumsum(mask, axis=0)
# Step 4: mask = [1, 1, 0, 0]
mask = tf.maximum(1. - mask, 0.)
# N.B. this reduces the length of weighted_auto_corr by a factor of 2.
# It still works fine in the formula below.
weighted_auto_corr = _sum_pairs(weighted_auto_corr) * mask
elif filter_threshold is not None:
filter_threshold = tf.convert_to_tensor(filter_threshold,
dtype=dt,
name='filter_threshold')
# Get a binary mask to zero out values of auto_corr below the threshold.
# mask[i, ...] = 1 if auto_corr[j, ...] > threshold for all j <= i,
# mask[i, ...] = 0, otherwise.
# So, along dimension zero, the mask will look like [1, 1, ..., 0, 0,...]
# Building step by step,
# Assume auto_corr = [1, 0.5, 0.0, 0.3], and filter_threshold = 0.2.
# Step 1: mask = [False, False, True, False]
mask = auto_corr < filter_threshold
# Step 2: mask = [0, 0, 1, 0]
mask = tf.cast(mask, dtype=dt)
# Step 3: mask = [0, 0, 1, 1]
mask = tf.cumsum(mask, axis=0)
# Step 4: mask = [1, 1, 0, 0]
mask = tf.maximum(1. - mask, 0.)
weighted_auto_corr *= mask
return num_chains * n / (-1 +
2 * tf.reduce_sum(weighted_auto_corr, axis=0))
def _parameter_control_dependencies(self, is_init):
assertions = []
# Check num_steps is a scalar that's at least 1.
if is_init != tensor_util.is_ref(self.num_steps):
num_steps = tf.convert_to_tensor(self.num_steps)
num_steps_ = tf.get_static_value(num_steps)
if num_steps_ is not None:
if np.ndim(num_steps_) != 0:
raise ValueError(
'`num_steps` must be a scalar but it has rank {}'.format(
np.ndim(num_steps_)))
if num_steps_ < 1:
raise ValueError('`num_steps` must be at least 1.')
elif self.validate_args:
message = '`num_steps` must be a scalar'
assertions.append(
assert_util.assert_rank_at_most(self.num_steps, 0, message=message))
assertions.append(
assert_util.assert_greater_equal(
num_steps, 1,
message='`num_steps` must be at least 1.'))
# Check that the initial distribution has scalar events over the
# integers.
if is_init and not dtype_util.is_integer(self.initial_distribution.dtype):
raise ValueError(
'`initial_distribution.dtype` ({}) is not over integers'.format(
dtype_util.name(self.initial_distribution.dtype)))
if tensorshape_util.rank(self.initial_distribution.event_shape) is not None:
if tensorshape_util.rank(self.initial_distribution.event_shape) != 0:
raise ValueError('`initial_distribution` must have scalar `event_dim`s')
elif self.validate_args:
assertions += [
assert_util.assert_equal(
ps.size(self.initial_distribution.event_shape_tensor()),
0,
message='`initial_distribution` must have scalar `event_dim`s'),
]
# Check that the transition distribution is over the integers.
if (is_init and
not dtype_util.is_integer(self.transition_distribution.dtype)):
raise ValueError(
'`transition_distribution.dtype` ({}) is not over integers'.format(
dtype_util.name(self.transition_distribution.dtype)))
# Check observations have non-scalar batches.
# The graph version of this assertion is incorporated as
# a control dependency of the transition/observation
# compatibility test.
if tensorshape_util.rank(self.observation_distribution.batch_shape) == 0:
raise ValueError(
"`observation_distribution` can't have scalar batches")
# Check transitions have non-scalar batches.
# The graph version of this assertion is incorporated as
# a control dependency of the transition/observation
# compatibility test.
if tensorshape_util.rank(self.transition_distribution.batch_shape) == 0:
raise ValueError(
"`transition_distribution` can't have scalar batches")
# Check compatibility of transition distribution and observation
# distribution.
tdbs = self.transition_distribution.batch_shape
odbs = self.observation_distribution.batch_shape
if (tensorshape_util.dims(tdbs) is not None and
tf.compat.dimension_value(odbs[-1]) is not None):
if (tf.compat.dimension_value(tdbs[-1]) !=
tf.compat.dimension_value(odbs[-1])):
raise ValueError(
'`transition_distribution` and `observation_distribution` '
'must agree on last dimension of batch size')
elif self.validate_args:
tdbs = self.transition_distribution.batch_shape_tensor()
odbs = self.observation_distribution.batch_shape_tensor()
transition_precondition = assert_util.assert_greater(
ps.size(tdbs), 0,
message=('`transition_distribution` can\'t have scalar '
'batches'))
observation_precondition = assert_util.assert_greater(
ps.size(odbs), 0,
message=('`observation_distribution` can\'t have scalar '
'batches'))
with tf.control_dependencies([
transition_precondition,
observation_precondition]):
assertions += [
assert_util.assert_equal(
tdbs[-1],
odbs[-1],
message=('`transition_distribution` and '
'`observation_distribution` '
'must agree on last dimension of batch size'))]
return assertions
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 _parameter_control_dependencies(self, is_init):
assertions = []
axis = None
paddings = None
if is_init != tensor_util.is_ref(self.axis):
# First we check the shape of the axis argument.
msg = 'Argument `axis` must be scalar or vector.'
if tensorshape_util.rank(self.axis.shape) is not None:
if tensorshape_util.rank(self.axis.shape) > 1:
raise ValueError(msg)
elif self.validate_args:
if axis is None: axis = tf.convert_to_tensor(self.axis)
assertions.append(assert_util.assert_rank_at_most(
axis, 1, message=msg))
# Next we check the values of the axis argument.
axis_ = tf.get_static_value(self.axis)
msg = 'Argument `axis` must be negative.'
if axis_ is not None:
if np.any(axis_ > -1):
raise ValueError(msg)
elif self.validate_args:
if axis is None: axis = tf.convert_to_tensor(self.axis)
assertions.append(assert_util.assert_less(axis, 0, message=msg))
msg = 'Argument `axis` elements must be unique.'
if axis_ is not None:
if len(np.array(axis_).reshape(-1)) != len(np.unique(axis_)):
raise ValueError(msg)
elif self.validate_args:
if axis is None: axis = tf.convert_to_tensor(self.axis)
assertions.append(assert_util.assert_equal(
prefer_static.size0(axis),
prefer_static.size0(prefer_static.setdiff1d(axis)),
message=msg))
if is_init != tensor_util.is_ref(self.paddings):
# First we check the shape of the paddings argument.
msg = 'Argument `paddings` must be a vector of pairs.'
if tensorshape_util.is_fully_defined(self.paddings.shape):
shape = np.int32(self.paddings.shape)
if len(shape) != 2 or shape[0] < 1 or shape[1] != 2:
raise ValueError(msg)
elif self.validate_args:
if paddings is None: paddings = tf.convert_to_tensor(self.paddings)
with tf.control_dependencies([
assert_util.assert_equal(tf.rank(paddings), 2, message=msg)]):
shape = tf.shape(paddings)
assertions.extend([
assert_util.assert_greater(shape[0], 0, message=msg),
assert_util.assert_equal(shape[1], 2, message=msg),
])
# Next we check the values of the paddings argument.
paddings_ = tf.get_static_value(self.paddings)
msg = 'Argument `paddings` must be non-negative.'
if paddings_ is not None:
if np.any(paddings_ < 0):
raise ValueError(msg)
elif self.validate_args:
if paddings is None: paddings = tf.convert_to_tensor(self.paddings)
assertions.append(assert_util.assert_greater(
paddings, -1, message=msg))
if is_init != (tensor_util.is_ref(self.axis) and
tensor_util.is_ref(self.paddings)):
axis_ = tf.get_static_value(self.axis)
if axis_ is None and axis is None:
axis = tf.convert_to_tensor(self.axis)
len_axis = prefer_static.size0(prefer_static.reshape(
axis if axis_ is None else axis_, shape=-1))
paddings_ = tf.get_static_value(self.paddings)
if paddings_ is None and paddings is None:
paddings = tf.convert_to_tensor(self.paddings)
len_paddings = prefer_static.size0(
paddings if paddings_ is None else paddings_)
msg = ('Arguments `axis` and `paddings` must have the same number '
'of elements.')
if (prefer_static.is_numpy(len_axis) and
prefer_static.is_numpy(len_paddings)):
if len_axis != len_paddings:
raise ValueError(msg + ' Saw: {}, {}.'.format(
self.axis, self.paddings))
elif self.validate_args:
assertions.append(assert_util.assert_equal(
len_axis, len_paddings, message=msg))
return assertions
def kendalls_tau(y_true, y_pred, name=None):
"""Computes Kendall's Tau for two ordered lists.
Kendall's Tau measures the correlation between ordinal rankings. This
implementation is similar to the one used in scipy.stats.kendalltau.
The provided values may be of any type that is sortable, with the
argsort indices indicating the true or proposed ordinal sequence.
Args:
y_true: a `Tensor` of shape `[n]` containing the true ordinal ranking.
y_pred: a `Tensor` of shape `[n]` containing the predicted ordering of the
same N items.
name: Optional Python `str` name for ops created by this method.
Default value: `None` (i.e., 'kendalls_tau').
Returns:
kendalls_tau: Kendall's Tau, the 1945 tau-b formulation that ignores
ordering of ties, as a `float32` scalar Tensor.
"""
with tf.name_scope(name or 'kendalls_tau'):
in_type = dtype_util.common_dtype([y_true, y_pred],
dtype_hint=tf.float32)
y_true = tf.convert_to_tensor(y_true, name='y_true', dtype=in_type)
y_pred = tf.convert_to_tensor(y_pred, name='y_pred', dtype=in_type)
tensorshape_util.assert_is_compatible_with(y_true.shape, y_pred.shape)
assertions = [
assert_util.assert_rank(y_true, 1),
assert_util.assert_greater(
ps.size(y_true), 1, 'Ordering requires at least 2 elements.')
]
with tf.control_dependencies(assertions):
lexa = lexicographical_indirect_sort(y_true, y_pred)
# See A Computer Method for Calculating Kendall's Tau with Ungrouped Data
# by William Night, Journal of the American Statistical Association,
# Jun., 1966, Vol. 61, No. 314, Part 1 (Jun., 1966), pp. 436-439
# for notation https://www.jstor.org/stable/2282833
def jointly_tied_pairs_body(first, t, i):
not_equal = tf.math.logical_or(
tf.not_equal(y_true[lexa[first]], y_true[lexa[i]]),
tf.not_equal(y_pred[lexa[first]], y_pred[lexa[i]]))
return (tf.where(not_equal, i, first),
tf.where(not_equal,
t + ((i - first) * (i - first - 1)) // 2,
t), i + 1)
n = ps.size0(y_true)
first, t, _ = tf.while_loop(cond=lambda first, t, i: i < n,
body=jointly_tied_pairs_body,
loop_vars=(0, 0, 1))
t += ((n - first) * (n - first - 1)) // 2
def ties_y_true_body(first, v, i):
not_equal = tf.not_equal(y_true[lexa[first]], y_true[lexa[i]])
return (tf.where(not_equal, i, first),
tf.where(not_equal,
v + ((i - first) * (i - first - 1)) // 2,
v), i + 1)
first, v, _ = tf.while_loop(cond=lambda first, v, i: i < n,
body=ties_y_true_body,
loop_vars=(0, 0, 1))
v += ((n - first) * (n - first - 1)) // 2
# count exchanges
exchanges, newperm = iterative_mergesort(y_pred, lexa)
def ties_in_y_pred_body(first, u, i):
not_equal = tf.not_equal(y_pred[newperm[first]],
y_pred[newperm[i]])
return (tf.where(not_equal, i, first),
tf.where(not_equal,
u + ((i - first) * (i - first - 1)) // 2,
u), i + 1)
first, u, _ = tf.while_loop(cond=lambda first, u, i: i < n,
body=ties_in_y_pred_body,
loop_vars=(0, 0, 1))
u += ((n - first) * (n - first - 1)) // 2
n0 = (n * (n - 1)) // 2
assertions = [
assert_util.assert_less(v, tf.cast(n0, tf.int32),
'All ranks are ties for y_true.'),
assert_util.assert_less(u, tf.cast(n0, tf.int32),
'All ranks are ties for y_pred.')
]
with tf.control_dependencies(assertions):
return (tf.cast(n0 - (u + v - t), tf.float32) -
2.0 * tf.cast(exchanges, tf.float32)) / tf.math.sqrt(
tf.cast(n0 - v, tf.float32) *
tf.cast(n0 - u, tf.float32))