def __init__(self,
component_ssms,
constant_offset=0.,
observation_noise_scale=None,
initial_state_prior=None,
initial_step=0,
validate_args=False,
allow_nan_stats=True,
name=None):
"""Build a state space model representing the sum of component models.
Args:
component_ssms: Python `list` containing one or more
`tfd.LinearGaussianStateSpaceModel` instances. The components
will in general implement different time-series models, with possibly
different `latent_size`, but they must have the same `dtype`, event
shape (`num_timesteps` and `observation_size`), and their batch shapes
must broadcast to a compatible batch shape.
constant_offset: `float` `Tensor` of shape broadcasting to
`concat([batch_shape, [num_timesteps]]`) specifying a constant value
added to the sum of outputs from the component models. This allows the
components to model the shifted series
`observed_time_series - constant_offset`.
Default value: `0.`
observation_noise_scale: Optional scalar `float` `Tensor` indicating the
standard deviation of the observation noise. May contain additional
batch dimensions, which must broadcast with the batch shape of elements
in `component_ssms`. If `observation_noise_scale` is specified for the
`AdditiveStateSpaceModel`, the observation noise scales of component
models are ignored. If `None`, the observation noise scale is derived
by summing the noise variances of the component models, i.e.,
`observation_noise_scale = sqrt(sum(
[ssm.observation_noise_scale**2 for ssm in component_ssms]))`.
initial_state_prior: Optional instance of `tfd.MultivariateNormal`
representing a prior distribution on the latent state at time
`initial_step`. If `None`, defaults to the independent priors from
component models, i.e.,
`[component.initial_state_prior for component in component_ssms]`.
Default value: `None`.
initial_step: Optional scalar `int` `Tensor` specifying the starting
timestep.
Default value: 0.
validate_args: Python `bool`. Whether to validate input
with asserts. If `validate_args` is `False`, and the inputs are
invalid, correct behavior is not guaranteed.
Default value: `False`.
allow_nan_stats: Python `bool`. If `False`, raise an
exception if a statistic (e.g. mean/mode/etc...) is undefined for any
batch member. If `True`, batch members with valid parameters leading to
undefined statistics will return NaN for this statistic.
Default value: `True`.
name: Python `str` name prefixed to ops created by this class.
Default value: "AdditiveStateSpaceModel".
Raises:
ValueError: if components have different `num_timesteps`.
"""
with tf.name_scope(name or 'AdditiveStateSpaceModel') as name:
# Check that all components have the same dtype
dtype = tf.debugging.assert_same_float_dtype(component_ssms)
# Convert scalar offsets to canonical shape `[..., num_timesteps]`.
constant_offset = (tf.convert_to_tensor(
value=constant_offset, name='constant_offset', dtype=dtype) *
tf.ones([1], dtype=dtype))
offset_length = prefer_static.shape(constant_offset)[-1]
assertions = []
# Construct an initial state prior as a block-diagonal combination
# of the component state priors.
if initial_state_prior is None:
initial_state_prior = sts_util.factored_joint_mvn(
[ssm.initial_state_prior for ssm in component_ssms])
dtype = initial_state_prior.dtype
static_num_timesteps = [
tf.get_static_value(ssm.num_timesteps)
for ssm in component_ssms
if tf.get_static_value(ssm.num_timesteps) is not None
]
# If any components have a static value for `num_timesteps`, use that
# value for the additive model. (and check that all other static values
# match it).
if static_num_timesteps:
num_timesteps = static_num_timesteps[0]
if not all([
component_timesteps == num_timesteps
for component_timesteps in static_num_timesteps
]):
raise ValueError(
'Additive model components must all have the same '
'number of timesteps '
'(saw: {})'.format(static_num_timesteps))
else:
num_timesteps = component_ssms[0].num_timesteps
if validate_args and len(static_num_timesteps) != len(
component_ssms):
assertions += [
tf.debugging.assert_equal( # pylint: disable=g-complex-comprehension
num_timesteps,
ssm.num_timesteps,
message='Additive model components must all have '
'the same number of timesteps.')
for ssm in component_ssms
]
# Define the transition and observation models for the additive SSM.
# See the "mathematical details" section of the class docstring for
# further information. Note that we define these as callables to
# handle the fully general case in which some components have time-
# varying dynamics.
def transition_matrix_fn(t):
return tfl.LinearOperatorBlockDiag([
ssm.get_transition_matrix_for_timestep(t)
for ssm in component_ssms
])
def transition_noise_fn(t):
return sts_util.factored_joint_mvn([
ssm.get_transition_noise_for_timestep(t)
for ssm in component_ssms
])
# Build the observation matrix, concatenating (broadcast) observation
# matrices from components. We also take this as an opportunity to enforce
# any dynamic assertions we may have generated above.
broadcast_batch_shape = tf.convert_to_tensor(
value=sts_util.broadcast_batch_shape([
ssm.get_observation_matrix_for_timestep(initial_step)
for ssm in component_ssms
]),
dtype=tf.int32)
broadcast_obs_matrix = tf.ones(tf.concat(
[broadcast_batch_shape, [1, 1]], axis=0),
dtype=dtype)
if assertions:
with tf.control_dependencies(assertions):
broadcast_obs_matrix = tf.identity(broadcast_obs_matrix)
def observation_matrix_fn(t):
return tfl.LinearOperatorFullMatrix(
tf.concat([
ssm.get_observation_matrix_for_timestep(t).to_dense() *
broadcast_obs_matrix for ssm in component_ssms
],
axis=-1))
# Broadcast the constant offset across timesteps.
offset_at_step = lambda t: ( # pylint: disable=g-long-lambda
constant_offset if offset_length == 1 else tf.gather(
constant_offset, tf.minimum(t, offset_length - 1), axis=-1)
[..., tf.newaxis])
if observation_noise_scale is not None:
observation_noise_scale = tf.convert_to_tensor(
value=observation_noise_scale,
name='observation_noise_scale',
dtype=dtype)
def observation_noise_fn(t):
return tfd.MultivariateNormalDiag(
loc=(sum([
ssm.get_observation_noise_for_timestep(t).mean()
for ssm in component_ssms
]) + offset_at_step(t)),
scale_diag=observation_noise_scale[..., tf.newaxis])
else:
def observation_noise_fn(t):
offset = offset_at_step(t)
return sts_util.sum_mvns([
tfd.MultivariateNormalDiag(
loc=offset, scale_diag=tf.zeros_like(offset))
] + [
ssm.get_observation_noise_for_timestep(t)
for ssm in component_ssms
])
super(AdditiveStateSpaceModel,
self).__init__(num_timesteps=num_timesteps,
transition_matrix=transition_matrix_fn,
transition_noise=transition_noise_fn,
observation_matrix=observation_matrix_fn,
observation_noise=observation_noise_fn,
initial_state_prior=initial_state_prior,
initial_step=initial_step,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
name=name)
def _cdf(self, x): loc = tf.identity(self.loc) return tf.cast(x >= loc - self._slack(loc), dtype=self.dtype)
def compute_alignment_loss(embs,
batch_size,
steps=None,
seq_lens=None,
stochastic_matching=False,
normalize_embeddings=False,
loss_type='classification',
similarity_type='l2',
num_cycles=20,
cycle_length=2,
temperature=0.1,
label_smoothing=0.1,
variance_lambda=0.001,
huber_delta=0.1,
normalize_indices=True):
"""Computes alignment loss between sequences of embeddings.
This function is a wrapper around different variants of the alignment loss
described deterministic_alignment.py and stochastic_alignment.py files. The
structure of the library is as follows:
i) loss_fns.py - Defines the different loss functions.
ii) deterministic_alignment.py - Performs the alignment between sequences by
deterministically sampling all steps of the sequences.
iii) stochastic_alignment.py - Performs the alignment between sequences by
stochasticallty sub-sampling a fixed number of steps from the sequences.
There are four major hparams that need to be tuned while applying the loss:
i) Should the loss be applied with L2 normalization on the embeddings or
without it?
ii) Should we perform stochastic alignment of sequences? This means should we
use all the steps of the embedding or only choose a random subset for
alignment?
iii) Should we apply cycle-consistency constraints using a classification loss
or a regression loss? (Section 3 in paper)
iv) Should the similarity metric be based on an L2 distance or cosine
similarity?
Other hparams that can be used to control how hard/soft we want the alignment
between different sequences to be:
i) temperature (all losses)
ii) label_smoothing (classification)
iii) variance_lambda (regression_mse_var)
iv) huber_delta (regression_huber)
Each of these params are used in their respective loss types (in brackets) and
allow the application of the cycle-consistency constraints in a controllable
manner but they do so in very different ways. Please refer to paper for more
details.
The default hparams work well for frame embeddings of videos of humans
performing actions. Other datasets might need different values of hparams.
Args:
embs: Tensor, sequential embeddings of the shape [N, T, D] where N is the
batch size, T is the number of timesteps in the sequence, D is the size of
the embeddings.
batch_size: Integer, Size of the batch.
steps: Tensor, step indices/frame indices of the embeddings of the shape
[N, T] where N is the batch size, T is the number of the timesteps.
If this is set to None, then we assume that the sampling was done in a
uniform way and use tf.range(num_steps) as the steps.
seq_lens: Tensor, Lengths of the sequences from which the sampling was done.
This can provide additional information to the alignment loss. This is
different from num_steps which is just the number of steps that have been
sampled from the entire sequence.
stochastic_matching: Boolean, Should the used for matching be sampled
stochastically or deterministically? Deterministic is better for TPU.
Stochastic is better for adding more randomness to the training process
and handling long sequences.
normalize_embeddings: Boolean, Should the embeddings be normalized or not?
Default is to use raw embeddings. Be careful if you are normalizing the
embeddings before calling this function.
loss_type: String, This specifies the kind of loss function to use.
Currently supported loss functions: classification, regression_mse,
regression_mse_var, regression_huber.
similarity_type: String, Currently supported similarity metrics: l2, cosine.
num_cycles: Integer, number of cycles to match while aligning
stochastically. Only used in the stochastic version.
cycle_length: Integer, Lengths of the cycle to use for matching. Only used
in the stochastic version. By default, this is set to 2.
temperature: Float, temperature scaling used to scale the similarity
distributions calculated using the softmax function.
label_smoothing: Float, Label smoothing argument used in
tf.keras.losses.categorical_crossentropy function and described in this
paper https://arxiv.org/pdf/1701.06548.pdf.
variance_lambda: Float, Weight of the variance of the similarity
predictions while cycling back. If this is high then the low variance
similarities are preferred by the loss while making this term low results
in high variance of the similarities (more uniform/random matching).
huber_delta: float, Huber delta described in tf.keras.losses.huber_loss.
normalize_indices: Boolean, If True, normalizes indices by sequence lengths.
Useful for ensuring numerical instabilities doesn't arise as sequence
indices can be large numbers.
Returns:
loss: Tensor, Scalar loss tensor that imposes the chosen variant of the
cycle-consistency loss.
"""
##############################################################################
# Checking inputs and setting defaults.
##############################################################################
# Get the number of timestemps in the sequence embeddings.
num_steps = tf.shape(embs)[1]
# If steps has not been provided assume sampling has been done uniformly.
if steps is None:
steps = tf.tile(tf.expand_dims(tf.range(num_steps), axis=0),
[batch_size, 1])
# If seq_lens has not been provided assume is equal to the size of the
# time axis in the emebeddings.
if seq_lens is None:
seq_lens = tf.tile(tf.expand_dims(num_steps, 0), [batch_size])
if not tf.executing_eagerly():
# Check if batch size embs is consistent with provided batch size.
with tf.control_dependencies(
[tf.assert_equal(batch_size,
tf.shape(embs)[0])]):
embs = tf.identity(embs)
# Check if number of timesteps in embs is consistent with provided steps.
with tf.control_dependencies([
tf.assert_equal(num_steps,
tf.shape(steps)[1]),
tf.assert_equal(batch_size,
tf.shape(steps)[0])
]):
steps = tf.identity(steps)
else:
tf.assert_equal(batch_size, tf.shape(steps)[0])
tf.assert_equal(num_steps, tf.shape(steps)[1])
tf.assert_equal(batch_size, tf.shape(embs)[0])
##############################################################################
# Perform alignment and return loss.
##############################################################################
if normalize_embeddings:
embs = tf.nn.l2_normalize(embs, axis=-1)
if stochastic_matching:
loss = compute_stochastic_alignment_loss(
embs=embs,
steps=steps,
seq_lens=seq_lens,
num_steps=num_steps,
batch_size=batch_size,
loss_type=loss_type,
similarity_type=similarity_type,
num_cycles=num_cycles,
cycle_length=cycle_length,
temperature=temperature,
label_smoothing=label_smoothing,
variance_lambda=variance_lambda,
huber_delta=huber_delta,
normalize_indices=normalize_indices)
else:
loss = compute_deterministic_alignment_loss(
embs=embs,
steps=steps,
seq_lens=seq_lens,
num_steps=num_steps,
batch_size=batch_size,
loss_type=loss_type,
similarity_type=similarity_type,
temperature=temperature,
label_smoothing=label_smoothing,
variance_lambda=variance_lambda,
huber_delta=huber_delta,
normalize_indices=normalize_indices)
return loss
def _forward_log_det_jacobian(self, x):
"""Returns the (log o det o Jacobian o g)(x)."""
return tf.identity(x)
def _batch_shape_tensor(self):
return tf.identity(self._batch_shape_val)
def logits_parameter(self, name=None):
"""Logits computed from non-`None` input arg (`probs` or `logits`)."""
with self._name_and_control_scope(name or 'logits_parameter'):
if self.logits is None:
return tf.math.log(self.probs) - tf.math.log1p(-self.probs)
return tf.identity(self.logits)
def __init__(self,
total_count,
logits=None,
probs=None,
validate_args=False,
allow_nan_stats=True,
name='NegativeBinomial'):
"""Construct NegativeBinomial distributions.
Args:
total_count: Non-negative floating-point `Tensor` with shape
broadcastable to `[B1,..., Bb]` with `b >= 0` and the same dtype as
`probs` or `logits`. Defines this as a batch of `N1 x ... x Nm`
different Negative Binomial distributions. In practice, this represents
the number of negative Bernoulli trials to stop at (the `total_count`
of failures), but this is still a valid distribution when
`total_count` is a non-integer.
logits: Floating-point `Tensor` with shape broadcastable to
`[B1, ..., Bb]` where `b >= 0` indicates the number of batch dimensions.
Each entry represents logits for the probability of success for
independent Negative Binomial distributions and must be in the open
interval `(-inf, inf)`. Only one of `logits` or `probs` should be
specified.
probs: Positive floating-point `Tensor` with shape broadcastable to
`[B1, ..., Bb]` where `b >= 0` indicates the number of batch dimensions.
Each entry represents the probability of success for independent
Negative Binomial distributions and must be in the open interval
`(0, 1)`. Only one of `logits` or `probs` should be specified.
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.
"""
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([total_count, logits, probs],
dtype_hint=tf.float32)
self._logits, self._probs = distribution_util.get_logits_and_probs(
logits,
probs,
validate_args=validate_args,
name=name,
dtype=dtype)
total_count = tf.convert_to_tensor(total_count,
name='total_count',
dtype=dtype)
with tf.control_dependencies(
[assert_util.assert_positive(total_count
)] if validate_args else []):
self._total_count = tf.identity(total_count,
name='total_count')
super(NegativeBinomial, self).__init__(
dtype=self._probs.dtype,
reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
parameters=parameters,
graph_parents=[self._total_count, self._probs, self._logits],
name=name)
def _event_shape_tensor(self):
with tf.control_dependencies(self._runtime_assertions):
return tf.identity(self.distribution.event_shape_tensor())
def _logits_parameter_no_checks(self):
if self._logits is None:
probs = tf.convert_to_tensor(self._probs)
return tf.math.log(probs) - tf.math.log1p(-probs)
return tf.identity(self._logits)
def _mean(self):
return tf.identity(self.probs)
def _batch_shape_tensor(self):
with tf.control_dependencies(self._runtime_assertions):
return tf.identity(self._batch_shape_)
def usage_counting_identity(var):
key = (id(var), var.name)
VAR_USAGES[key] = VAR_USAGES.get(key,
[]) + [traceback.format_stack(limit=15)]
return tf.identity(var)
def _probs_parameter_no_checks(self):
if self._logits is None:
return tf.identity(self._probs)
return tf.math.softmax(self._logits)
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 _single_identity_op_at_end():
inputs = keras.Input(shape=(10, ))
x = keras.layers.Dense(10)(inputs)
outputs = tf.identity(x)
return keras.Model(inputs, outputs)
def prevent_constant_folding(self, tensor, dummy_inputs):
tensor = tf.identity(tensor)
outputs = tf.cond(tf.reduce_sum(dummy_inputs) > 0,
lambda: tensor,
lambda: tensor)
return outputs
def test_identity(self):
x = keras.Input(shape=(1, ))
ident = tf.identity(x)
# This is now a graph tensor, and should be able to continue in graphland
self.assertIn('Identity', ident.name)
def xform_static(x):
# Copy the Tensor, because otherwise the set_shape can pass information
# into the past.
x = tf.identity(x)
tensorshape_util.set_shape(x, [1])
return x
def probs_parameter(self, name=None):
"""Probs computed from non-`None` input arg (`probs` or `logits`)."""
with self._name_and_control_scope(name or 'probs_parameter'):
if self.logits is None:
return tf.identity(self.probs)
return tf.math.sigmoid(self.logits)
def forward_event_shape_tensor(self,
input_shape,
name='forward_event_shape_tensor'):
"""Shape of a sample from a single batch as a list of `int32` 1D `Tensor`s.
Args:
input_shape: `Tensor`, `int32` vector indicating event-portion shape
passed into `forward` function.
name: name to give to the op
Returns:
forward_event_shape_tensor: A list of `Tensor`, `int32` vectors indicating
event-portion shape after applying `forward`. The list is of the same
length as `split_sizes`.
"""
with self._name_and_control_scope(name):
input_shape = tf.convert_to_tensor(input_shape,
dtype_hint=tf.int32,
name='input_shape')
# Validate `input_shape` statically if possible and get assertions.
is_validated = self._validate_input_shape(
tensorshape_util.constant_value_as_shape(input_shape))
if is_validated or not self.validate_args:
assertions = []
else:
assertions = self._validate_input_shape_tensor(input_shape)
with tf.control_dependencies(assertions):
if self.split_sizes is None:
split_sizes = tf.convert_to_tensor(
[input_shape[self.axis] // self.num_splits] *
self.num_splits)
else:
# Deduce the value of the unknown element of `split_sizes`, if any.
split_sizes = tf.convert_to_tensor(self.split_sizes)
split_sizes = tf.where(
split_sizes < 0,
input_shape[self.axis] - tf.reduce_sum(split_sizes) -
1, # Cancel the unknown size `-1`.
split_sizes)
# Each element of the `output_shape_tensor` list is equal to the
# `input_shape`, with the corresponding element of `split_sizes`
# substituted in the `axis` position.
positive_axis = prefer_static.rank_from_shape(
input_shape) + self.axis
tiled_input_shape = tf.tile(input_shape[tf.newaxis, :],
[self.num_splits, 1])
fused_output_shapes = tf.concat([
tiled_input_shape[:, :positive_axis],
split_sizes[..., tf.newaxis],
tiled_input_shape[:, positive_axis + 1:]
],
axis=1)
output_shapes = tf.unstack(fused_output_shapes,
num=self.num_splits)
return [
tf.identity(
tf.convert_to_tensor(t,
dtype_hint=tf.int32,
name='forward_event_shape'))
for t in output_shapes
]
def __call__(self, x):
assert x.dtype == tf.float32
return tf.identity(x)
def wrapper(*args):
"""Wrapper that wraps/unwraps args, retvals, and runs the function."""
if _pmap_config.devices() is not None:
raise ValueError(
"Found a surrounding pmap. Nested pmap is not supported "
"yet.")
# TODO(wangpeng): Maybe we should use `asarray` to convert everything to
# ndarray first.
args = _np_to_tf(args)
flattened_input_args = tf.nest.flatten(args)
flattened_per_device_args = [[] for _ in devices]
for arg in flattened_input_args:
if isinstance(arg, tf.Tensor):
# TODO(nareshmodi): Try and use the dynamic shape instead.
if (not arg.shape.rank) or arg.shape[0] != len(devices):
# TODO(nareshmodi): Fix this restriction
raise ValueError(
"Input tensors need to have a first dimension equal to "
"the number of devices; got tensor of shape %s and %s devices"
% (arg.shape, len(devices)))
# NOTE: Alternatively use tf.split, and place the split tensors on the
# appropriate device. The best solution for this is to have an API that
# splits a tensor across devices.
for j, device in enumerate(devices):
updated_arg = tf.gather(arg, j)
# TODO(wangpeng): Investigate whether we need a tf.identity for TPU.
if not has_tpu:
with tf.device(device):
updated_arg = tf.identity(updated_arg)
flattened_per_device_args[j].append(updated_arg)
elif isinstance(arg, ShardedNdArray):
for device_args, tensor in zip(flattened_per_device_args,
arg.tensors):
device_args.append(tensor)
else:
for device_args in flattened_per_device_args:
device_args.append(arg)
all_per_device_args = [
tf.nest.pack_sequence_as(args, device_args)
for device_args in flattened_per_device_args
]
with pmap_config(axis_name, devices):
results = pmap_fn(all_per_device_args)
# Rewrap things. This can probably be written better.
flattened_results = [tf.nest.flatten(result) for result in results]
final_tree = []
# TODO(nareshmodi): assert all items in flattened_results have the same
# structures
for i in range(len(flattened_results[0])):
tensors = []
for j, device in enumerate(devices):
assert isinstance(flattened_results[j][i], tf.Tensor), (
"currently only tensor return items are supported")
tensors.append(flattened_results[j][i])
final_tree.append(ShardedNdArray(tensors))
final_actual_result = tf.nest.pack_sequence_as(results[0], final_tree)
# Workaround b/121383831
if (has_tpu and isinstance(final_actual_result, list)
and len(final_actual_result)
== 1) and not _orig_result_is_list.val:
return final_actual_result[0]
else:
return final_actual_result
def preprocess(self, inputs):
true_image_shapes = [] # Doesn't matter for the fake model.
return tf.identity(inputs), true_image_shapes
def __init__(self,
concentration,
scale=None,
validate_args=False,
allow_nan_stats=True,
rate=None,
name="InverseGamma"):
"""Construct InverseGamma with `concentration` and `scale` parameters.
The parameters `concentration` and `scale` must be shaped in a way that
supports broadcasting (e.g. `concentration + scale` is a valid operation).
Args:
concentration: Floating point tensor, the concentration params of the
distribution(s). Must contain only positive values.
scale: Floating point tensor, the scale params of the distribution(s).
Must contain only positive values.
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.
rate: Deprecated (mis-named) alias for `scale`.
name: Python `str` name prefixed to Ops created by this class.
Raises:
TypeError: if `concentration` and `scale` are different dtypes.
"""
if rate is not None:
scale = rate
parameters = dict(locals())
with tf.name_scope(name) as name:
dtype = dtype_util.common_dtype([concentration, scale],
dtype_hint=tf.float32)
concentration = tf.convert_to_tensor(value=concentration,
name="concentration",
dtype=dtype)
scale = tf.convert_to_tensor(value=scale,
name="scale",
dtype=dtype)
with tf.control_dependencies([
assert_util.
assert_positive(concentration,
message="Concentration must be positive."),
assert_util.
assert_positive(scale, message="Scale must be positive."),
] if validate_args else []):
self._concentration = tf.identity(concentration,
name="concentration")
self._scale = tf.identity(scale, name="scale")
dtype_util.assert_same_float_dtype(
[self._concentration, self._scale])
super(InverseGamma, self).__init__(
dtype=self._concentration.dtype,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
parameters=parameters,
graph_parents=[self._concentration, self._scale],
name=name)
def _mean(self): return tf.identity(self.loc)
def test_posterior_marginals_high_rank(self, rank_o, rank_t, rank_i,
rank_s):
def increase_rank(n, x):
# By choosing prime number dimensions we make it less
# likely that a test will pass for accidental reasons.
primes = [3, 5, 7]
for i in range(n):
x = primes[i] * [x]
return x
observation_locs_data = tf.identity(
increase_rank(rank_o, tf.eye(4, dtype=self.dtype)))
observation_scales_data = tf.constant([0.25, 0.25, 0.25, 0.25],
dtype=self.dtype)
transition_matrix_data = tf.constant(increase_rank(
rank_t, [[0.8, 0.1, 0.1, 0.0], [0.1, 0.8, 0.0, 0.1],
[0.1, 0.0, 0.8, 0.1], [0.0, 0.1, 0.1, 0.8]]),
dtype=self.dtype)
initial_prob_data = tf.constant(increase_rank(
rank_i, [0.25, 0.25, 0.25, 0.25]),
dtype=self.dtype)
(initial_prob, transition_matrix, observation_locs,
observation_scales) = self.make_placeholders([
initial_prob_data, transition_matrix_data, observation_locs_data,
observation_scales_data
])
observations = tf.constant(increase_rank(
rank_s,
[[[0.91, 0.11], [0.21, 0.09]], [[0.11, 0.97], [0.12, 0.08]],
[[0.01, 0.12], [0.92, 0.11]], [[0.02, 0.11], [0.77, 0.11]],
[[0.81, 0.15], [0.21, 0.03]], [[0.01, 0.13], [0.23, 0.91]],
[[0.11, 0.12], [0.23, 0.79]], [[0.13, 0.11], [0.91, 0.29]]]),
dtype=self.dtype)
observation_distribution = tfp.distributions.TransformedDistribution(
tfd.MultivariateNormalDiag(observation_locs,
scale_diag=observation_scales),
tfp.bijectors.Reshape((2, 2)))
[num_steps] = self.make_placeholders([8])
model = tfd.HiddenMarkovModel(tfd.Categorical(probs=initial_prob),
tfd.Categorical(probs=transition_matrix),
observation_distribution,
num_steps=num_steps,
validate_args=True)
inferred_probs = self.evaluate(
model.posterior_marginals(observations).probs_parameter())
rank_e = max(rank_o, rank_t, rank_i, rank_s)
expected_probs = increase_rank(rank_e,
[[0.99994, 0.00000, 0.00006, 0.00000],
[0.45137, 0.01888, 0.52975, 0.00000],
[0.00317, 0.00002, 0.98112, 0.01570],
[0.00000, 0.00001, 0.99998, 0.00001],
[0.00495, 0.00001, 0.94214, 0.05289],
[0.00000, 0.00083, 0.00414, 0.99503],
[0.00000, 0.00000, 0.00016, 0.99984],
[0.00000, 0.00000, 0.99960, 0.00039]])
self.assertAllClose(inferred_probs, expected_probs, rtol=0., atol=1e-4)
def _validate_arg_if_not_none(arg, assertion, validate_args):
if arg is None:
return arg
with tf.control_dependencies([assertion(arg)] if validate_args else []):
result = tf.identity(arg)
return result
def call(self, inputs, mask=None): return tf.identity(inputs)
def _kl(a, b, name=None):
return tf.identity([float("nan")])
def call(self, inputs, training):
layers = []
for layer in self.layers:
inputs, block_layers = layer(inputs, training=training)
layers.append(block_layers)
return tf.identity(inputs, self._name), layers