def _inverse(self, y):
ndims = prefer_static.rank(y)
shifted_y = tf.pad(
tf.slice(
y, tf.zeros(ndims, dtype=tf.int32),
prefer_static.shape(y) -
tf.one_hot(ndims + self.axis, ndims, dtype=tf.int32)
), # Remove the last entry of y in the chosen dimension.
paddings=tf.one_hot(
tf.one_hot(ndims + self.axis, ndims, on_value=0, off_value=-1),
2,
dtype=tf.int32
) # Insert zeros at the beginning of the chosen dimension.
)
return y - shifted_y
def _mode(self):
logits = self._logits_parameter_no_checks()
ret = tf.one_hot(tf.argmax(logits, axis=-1),
self._event_size(logits),
dtype=self.dtype)
tensorshape_util.set_shape(ret, logits.shape)
return ret
def backward_step(most_likely_successor,
most_likely_given_successor):
return tf.reduce_sum(
input_tensor=(most_likely_given_successor *
tf.one_hot(most_likely_successor,
self._num_states,
dtype=tf.int64)),
axis=-1)
def _sample_one_batch_member(args):
logits, num_cat_samples = args[0], args[1] # [K], []
# x has shape [1, num_cat_samples = num_samples * num_trials]
x = tf.random.categorical(logits[tf.newaxis, ...],
num_cat_samples,
seed=seed)
x = tf.reshape(x, shape=[num_samples,
-1]) # [num_samples, num_trials]
x = tf.one_hot(
x, depth=num_classes) # [num_samples, num_trials, num_classes]
x = tf.reduce_sum(x, axis=-2) # [num_samples, num_classes]
return tf.cast(x, dtype=dtype)
def _sample_n(self, n, seed=None):
logits = self._logits_parameter_no_checks()
sample_shape = prefer_static.concat(
[[n], prefer_static.shape(logits)], 0)
event_size = self._event_size(logits)
if tensorshape_util.rank(logits.shape) == 2:
logits_2d = logits
else:
logits_2d = tf.reshape(logits, [-1, event_size])
samples = tf.random.categorical(logits_2d, n, seed=seed)
samples = tf.transpose(a=samples)
samples = tf.one_hot(samples, event_size, dtype=self.dtype)
ret = tf.reshape(samples, sample_shape)
return ret
def generate_step(state, _):
"""Take a single step in Markov chain."""
gen = self._transition_distribution.sample(n *
transition_repeat,
seed=strm())
# gen :: (n * transition_repeat) transition_batch
new_states = tf.reshape(gen, [n, batch_size, num_states])
# new_states :: n batch_size num_states
old_states_one_hot = tf.one_hot(state,
num_states,
dtype=tf.int32)
# old_states :: n batch_size num_states
return tf.reduce_sum(old_states_one_hot * new_states, axis=-1)
def _sample_n(self, n, seed):
with tf.control_dependencies(self._runtime_assertions):
seed = SeedStream(seed, salt="MixtureSameFamily")
x = self.components_distribution.sample(
n, seed=seed()) # [n, B, k, E]
# TODO(jvdillon): Consider using tf.gather (by way of index unrolling).
npdt = dtype_util.as_numpy_dtype(x.dtype)
mask = tf.one_hot(
indices=self.mixture_distribution.sample(
n, seed=seed()), # [n, B]
depth=self._num_components, # == k
on_value=npdt(1),
off_value=npdt(0)) # [n, B, k]
mask = distribution_utils.pad_mixture_dimensions(
mask, self, self.mixture_distribution,
self._event_ndims) # [n, B, k, [1]*e]
x = tf.reduce_sum(x * mask,
axis=-1 - self._event_ndims) # [n, B, E]
if self._reparameterize:
x = self._reparameterize_sample(x)
return x
def _sample_n(self, n, seed=None):
if self._use_static_graph:
with tf.control_dependencies(self._assertions):
# This sampling approach is almost the same as the approach used by
# `MixtureSameFamily`. The differences are due to having a list of
# `Distribution` objects rather than a single object, and maintaining
# random seed management that is consistent with the non-static code
# path.
samples = []
cat_samples = self.cat.sample(n, seed=seed)
stream = SeedStream(seed, salt="Mixture")
for c in range(self.num_components):
samples.append(self.components[c].sample(n, seed=stream()))
stack_axis = -1 - tensorshape_util.rank(self._static_event_shape)
x = tf.stack(samples, axis=stack_axis) # [n, B, k, E]
npdt = dtype_util.as_numpy_dtype(x.dtype)
mask = tf.one_hot(
indices=cat_samples, # [n, B]
depth=self._num_components, # == k
on_value=npdt(1),
off_value=npdt(0)) # [n, B, k]
mask = distribution_util.pad_mixture_dimensions(
mask, self, self._cat,
tensorshape_util.rank(self._static_event_shape)) # [n, B, k, [1]*e]
return tf.reduce_sum(x * mask, axis=stack_axis) # [n, B, E]
with tf.control_dependencies(self._assertions):
n = tf.convert_to_tensor(n, name="n")
static_n = tf.get_static_value(n)
n = int(static_n) if static_n is not None else n
cat_samples = self.cat.sample(n, seed=seed)
static_samples_shape = cat_samples.shape
if tensorshape_util.is_fully_defined(static_samples_shape):
samples_shape = tensorshape_util.as_list(static_samples_shape)
samples_size = tensorshape_util.num_elements(static_samples_shape)
else:
samples_shape = tf.shape(cat_samples)
samples_size = tf.size(cat_samples)
static_batch_shape = self.batch_shape
if tensorshape_util.is_fully_defined(static_batch_shape):
batch_shape = tensorshape_util.as_list(static_batch_shape)
batch_size = tensorshape_util.num_elements(static_batch_shape)
else:
batch_shape = self.batch_shape_tensor()
batch_size = tf.reduce_prod(batch_shape)
static_event_shape = self.event_shape
if tensorshape_util.is_fully_defined(static_event_shape):
event_shape = np.array(
tensorshape_util.as_list(static_event_shape), dtype=np.int32)
else:
event_shape = self.event_shape_tensor()
# Get indices into the raw cat sampling tensor. We will
# need these to stitch sample values back out after sampling
# within the component partitions.
samples_raw_indices = tf.reshape(tf.range(0, samples_size), samples_shape)
# Partition the raw indices so that we can use
# dynamic_stitch later to reconstruct the samples from the
# known partitions.
partitioned_samples_indices = tf.dynamic_partition(
data=samples_raw_indices,
partitions=cat_samples,
num_partitions=self.num_components)
# Copy the batch indices n times, as we will need to know
# these to pull out the appropriate rows within the
# component partitions.
batch_raw_indices = tf.reshape(
tf.tile(tf.range(0, batch_size), [n]), samples_shape)
# Explanation of the dynamic partitioning below:
# batch indices are i.e., [0, 1, 0, 1, 0, 1]
# Suppose partitions are:
# [1 1 0 0 1 1]
# After partitioning, batch indices are cut as:
# [batch_indices[x] for x in 2, 3]
# [batch_indices[x] for x in 0, 1, 4, 5]
# i.e.
# [1 1] and [0 0 0 0]
# Now we sample n=2 from part 0 and n=4 from part 1.
# For part 0 we want samples from batch entries 1, 1 (samples 0, 1),
# and for part 1 we want samples from batch entries 0, 0, 0, 0
# (samples 0, 1, 2, 3).
partitioned_batch_indices = tf.dynamic_partition(
data=batch_raw_indices,
partitions=cat_samples,
num_partitions=self.num_components)
samples_class = [None for _ in range(self.num_components)]
stream = SeedStream(seed, salt="Mixture")
for c in range(self.num_components):
n_class = tf.size(partitioned_samples_indices[c])
samples_class_c = self.components[c].sample(
n_class, seed=stream())
# Pull out the correct batch entries from each index.
# To do this, we may have to flatten the batch shape.
# For sample s, batch element b of component c, we get the
# partitioned batch indices from
# partitioned_batch_indices[c]; and shift each element by
# the sample index. The final lookup can be thought of as
# a matrix gather along locations (s, b) in
# samples_class_c where the n_class rows correspond to
# samples within this component and the batch_size columns
# correspond to batch elements within the component.
#
# Thus the lookup index is
# lookup[c, i] = batch_size * s[i] + b[c, i]
# for i = 0 ... n_class[c] - 1.
lookup_partitioned_batch_indices = (
batch_size * tf.range(n_class) + partitioned_batch_indices[c])
samples_class_c = tf.reshape(
samples_class_c, tf.concat([[n_class * batch_size], event_shape],
0))
samples_class_c = tf.gather(
samples_class_c,
lookup_partitioned_batch_indices,
name="samples_class_c_gather")
samples_class[c] = samples_class_c
# Stitch back together the samples across the components.
lhs_flat_ret = tf.dynamic_stitch(
indices=partitioned_samples_indices, data=samples_class)
# Reshape back to proper sample, batch, and event shape.
ret = tf.reshape(
lhs_flat_ret, tf.concat(
[samples_shape, self.event_shape_tensor()], 0))
tensorshape_util.set_shape(
ret,
tensorshape_util.concatenate(static_samples_shape, self.event_shape))
return ret
def _sample_n(self, n, seed=None):
with tf.control_dependencies(self._runtime_assertions):
strm = SeedStream(seed, salt="HiddenMarkovModel")
num_states = self._num_states
batch_shape = self.batch_shape_tensor()
batch_size = tf.reduce_prod(batch_shape)
# The batch sizes of the underlying initial distributions and
# transition distributions might not match the batch size of
# the HMM distribution.
# As a result we need to ask for more samples from the
# underlying distributions and then reshape the results into
# the correct batch size for the HMM.
init_repeat = (
tf.reduce_prod(self.batch_shape_tensor()) // tf.reduce_prod(
self._initial_distribution.batch_shape_tensor()))
init_state = self._initial_distribution.sample(n * init_repeat,
seed=strm())
init_state = tf.reshape(init_state, [n, batch_size])
# init_state :: n batch_size
transition_repeat = (
tf.reduce_prod(self.batch_shape_tensor()) // tf.reduce_prod(
self._transition_distribution.batch_shape_tensor()[:-1]))
def generate_step(state, _):
"""Take a single step in Markov chain."""
gen = self._transition_distribution.sample(n *
transition_repeat,
seed=strm())
# gen :: (n * transition_repeat) transition_batch
new_states = tf.reshape(gen, [n, batch_size, num_states])
# new_states :: n batch_size num_states
old_states_one_hot = tf.one_hot(state,
num_states,
dtype=tf.int32)
# old_states :: n batch_size num_states
return tf.reduce_sum(old_states_one_hot * new_states, axis=-1)
def _scan_multiple_steps():
"""Take multiple steps with tf.scan."""
dummy_index = tf.zeros(self._num_steps - 1, dtype=tf.float32)
if seed is not None:
# Force parallel_iterations to 1 to ensure reproducibility
# b/139210489
hidden_states = tf.scan(generate_step,
dummy_index,
initializer=init_state,
parallel_iterations=1)
else:
# Invoke default parallel_iterations behavior
hidden_states = tf.scan(generate_step,
dummy_index,
initializer=init_state)
# TODO(b/115618503): add/use prepend_initializer to tf.scan
return tf.concat([[init_state], hidden_states], axis=0)
hidden_states = prefer_static.cond(
self._num_steps > 1, _scan_multiple_steps,
lambda: init_state[tf.newaxis, ...])
hidden_one_hot = tf.one_hot(
hidden_states,
num_states,
dtype=self._observation_distribution.dtype)
# hidden_one_hot :: num_steps n batch_size num_states
# The observation distribution batch size might not match
# the required batch size so as with the initial and
# transition distributions we generate more samples and
# reshape.
observation_repeat = (batch_size // tf.reduce_prod(
self._observation_distribution.batch_shape_tensor()[:-1]))
possible_observations = self._observation_distribution.sample(
[self._num_steps, observation_repeat * n], seed=strm())
inner_shape = self._observation_distribution.event_shape
# possible_observations :: num_steps (observation_repeat * n)
# observation_batch[:-1] num_states inner_shape
possible_observations = tf.reshape(
possible_observations,
tf.concat([[self._num_steps, n], batch_shape, [num_states],
inner_shape],
axis=0))
# possible_observations :: steps n batch_size num_states inner_shape
hidden_one_hot = tf.reshape(
hidden_one_hot,
tf.concat([[self._num_steps, n], batch_shape, [num_states],
tf.ones_like(inner_shape)],
axis=0))
# hidden_one_hot :: steps n batch_size num_states "inner_shape"
observations = tf.reduce_sum(hidden_one_hot *
possible_observations,
axis=-1 - tf.size(inner_shape))
# observations :: steps n batch_size inner_shape
observations = distribution_util.move_dimension(
observations, 0, 1 + tf.size(batch_shape))
# returned :: n batch_shape steps inner_shape
return observations