def _build_test_model(self,
num_timesteps=5,
num_features=2,
batch_shape=(),
missing_prob=0,
true_noise_scale=0.1,
true_level_scale=0.04,
true_slope_scale=0.02,
prior_class=tfd.InverseGamma,
dtype=tf.float32):
seed = test_util.test_seed(sampler_type='stateless')
(design_seed, weights_seed, noise_seed, level_seed, slope_seed,
is_missing_seed) = samplers.split_seed(seed,
6,
salt='_build_test_model')
design_matrix = samplers.normal([num_timesteps, num_features],
dtype=dtype,
seed=design_seed)
weights = samplers.normal(list(batch_shape) + [num_features],
dtype=dtype,
seed=weights_seed)
regression = tf.linalg.matvec(design_matrix, weights)
noise = samplers.normal(list(batch_shape) + [num_timesteps],
dtype=dtype,
seed=noise_seed) * true_noise_scale
level_residuals = samplers.normal(list(batch_shape) + [num_timesteps],
dtype=dtype,
seed=level_seed) * true_level_scale
if true_slope_scale is not None:
slope = tf.cumsum(
samplers.normal(list(batch_shape) + [num_timesteps],
dtype=dtype,
seed=slope_seed) * true_slope_scale,
axis=-1)
level_residuals += slope
level = tf.cumsum(level_residuals, axis=-1)
time_series = (regression + noise + level)
is_missing = samplers.uniform(list(batch_shape) + [num_timesteps],
dtype=dtype,
seed=is_missing_seed) < missing_prob
model = gibbs_sampler.build_model_for_gibbs_fitting(
observed_time_series=tfp.sts.MaskedTimeSeries(
time_series[..., tf.newaxis], is_missing),
design_matrix=design_matrix,
weights_prior=tfd.Normal(loc=tf.cast(0., dtype),
scale=tf.cast(10.0, dtype)),
level_variance_prior=prior_class(concentration=tf.cast(
0.01, dtype),
scale=tf.cast(0.01 * 0.01,
dtype)),
slope_variance_prior=None if true_slope_scale is None else
prior_class(concentration=tf.cast(0.01, dtype),
scale=tf.cast(0.01 * 0.01, dtype)),
observation_noise_variance_prior=prior_class(
concentration=tf.cast(0.01, dtype),
scale=tf.cast(0.01 * 0.01, dtype)))
return model, time_series, is_missing
def _lower_triangular_mask(shape):
"""Creates a lower-triangular boolean mask over the last 2 dimensions."""
row_index = tf.cumsum(
tf.ones(shape=shape, dtype=tf.int32), axis=-2)
col_index = tf.cumsum(
tf.ones(shape=shape, dtype=tf.int32), axis=-1)
return tf.greater_equal(row_index, col_index)
def wasserstein_distance(u_values, v_values, u_weights, v_weights, p=1.0):
"""Differentiable 1-D Wasserstein distance.
Adapted from the scipy.stats implementation.
Args:
u_values: Samples from distribution `u`. Shape [batch_shape, n_samples].
v_values: Samples from distribution `v`. Shape [batch_shape, n_samples].
u_weights: Sample weights. Shape [batch_shape, n_samples].
v_weights: Sample weights. Shape [batch_shape, n_samples].
p: Degree of the distance norm. Wasserstein=1, Energy=2.
Returns:
The Wasserstein distance between samples. Shape [batch_shape].
"""
u_sorter = tf.argsort(u_values, axis=-1)
v_sorter = tf.argsort(v_values, axis=-1)
all_values = tf.concat([u_values, v_values], axis=-1)
all_values = tf.sort(all_values, axis=-1)
# Compute the differences between pairs of successive values of u and v.
deltas = spectral_ops.diff(all_values, axis=-1)
# Get the respective positions of the values of u and v among the values of
# both distributions.
batch_dims = len(u_values.shape) - 1
gather = lambda x, i: tf.gather(x, i, axis=-1, batch_dims=batch_dims)
u_cdf_indices = tf.searchsorted(gather(u_values, u_sorter),
all_values[..., :-1],
side='right')
v_cdf_indices = tf.searchsorted(gather(v_values, v_sorter),
all_values[..., :-1],
side='right')
# Calculate the CDFs of u and v using their weights, if specified.
if u_weights is None:
u_cdf = u_cdf_indices / float(u_values.shape[-1])
else:
u_sorted_cumweights = tf.concat([
tf.zeros_like(u_weights)[..., 0:1],
tf.cumsum(gather(u_weights, u_sorter), axis=-1)
],
axis=-1)
u_cdf = gather(u_sorted_cumweights, u_cdf_indices)
safe_divide(u_cdf, u_sorted_cumweights[..., -1:])
if v_weights is None:
v_cdf = v_cdf_indices / float(v_values.shape[-1])
else:
v_sorted_cumweights = tf.concat([
tf.zeros_like(v_weights)[..., 0:1],
tf.cumsum(gather(v_weights, v_sorter), axis=-1)
],
axis=-1)
v_cdf = gather(v_sorted_cumweights, v_cdf_indices)
safe_divide(v_cdf, v_sorted_cumweights[..., -1:])
# Compute the value of the integral based on the CDFs.
return tf.reduce_sum(deltas * tf.abs(u_cdf - v_cdf)**p, axis=-1)**(1.0 / p)
def get_dense_is_inside_for_dense_spans(
dense_start_positions: tf.Tensor,
dense_end_positions: tf.Tensor) -> tf.Tensor:
"""Dense mask whether position is inside span given dense starts / ends."""
# `tf.cumsum(dense_start_positions)[i]` computes how many spans start before
# or on the i-th position.
# `tf.cumsum(dense_end_positions, exclusive=True`) computes how many spans
# ends strictly before i-th positions.
# Their difference is equal to how many spans start before i-th position and
# end after or on i-th position. This is precisely how many spans contain
# i-th position.
is_inside_span = (tf.cumsum(dense_start_positions) -
tf.cumsum(dense_end_positions, exclusive=True))
# Adjust for the case of overlapping spans
is_inside_span = tf.minimum(is_inside_span, 1)
return is_inside_span
def _cdf(self, k):
# TODO(b/135263541): Improve numerical precision of categorical.cdf.
probs = self.probs_parameter()
num_categories = self._num_categories(probs)
k, probs = _broadcast_cat_event_and_params(
k, probs, base_dtype=dtype_util.base_dtype(self.dtype))
# Since the lowest number in the support is 0, any k < 0 should be zero in
# the output.
should_be_zero = k < 0
# Will use k as an index in the gather below, so clip it to {0,...,K-1}.
k = tf.clip_by_value(tf.cast(k, tf.int32), 0, num_categories - 1)
batch_shape = tf.shape(k)
# tf.gather(..., batch_dims=batch_dims) requires static batch_dims kwarg, so
# to handle the case where the batch shape is dynamic, flatten the batch
# dims (so we know batch_dims=1).
k_flat_batch = tf.reshape(k, [-1])
probs_flat_batch = tf.reshape(
probs, tf.concat(([-1], [num_categories]), axis=0))
cdf_flat = tf.gather(tf.cumsum(probs_flat_batch, axis=-1),
k_flat_batch[..., tf.newaxis],
batch_dims=1)
cdf = tf.reshape(cdf_flat, shape=batch_shape)
zero = np.array(0, dtype=dtype_util.as_numpy_dtype(cdf.dtype))
return tf.where(should_be_zero, zero, cdf)
def boolean_mask(boxlist, indicator, fields=None, scope=None,
use_static_shapes=False, indicator_sum=None):
"""Select boxes from BoxList according to indicator and return new BoxList.
`boolean_mask` returns the subset of boxes that are marked as "True" by the
indicator tensor. By default, `boolean_mask` returns boxes corresponding to
the input index list, as well as all additional fields stored in the boxlist
(indexing into the first dimension). However one can optionally only draw
from a subset of fields.
Args:
boxlist: BoxList holding N boxes
indicator: a rank-1 boolean tensor
fields: (optional) list of fields to also gather from. If None (default),
all fields are gathered from. Pass an empty fields list to only gather
the box coordinates.
scope: name scope.
use_static_shapes: Whether to use an implementation with static shape
gurantees.
indicator_sum: An integer containing the sum of `indicator` vector. Only
required if `use_static_shape` is True.
Returns:
subboxlist: a BoxList corresponding to the subset of the input BoxList
specified by indicator
Raises:
ValueError: if `indicator` is not a rank-1 boolean tensor.
"""
with tf.name_scope(scope, 'BooleanMask'):
if indicator.shape.ndims != 1:
raise ValueError('indicator should have rank 1')
if indicator.dtype != tf.bool:
raise ValueError('indicator should be a boolean tensor')
if use_static_shapes:
if not (indicator_sum and isinstance(indicator_sum, int)):
raise ValueError('`indicator_sum` must be a of type int')
selected_positions = tf.cast(indicator, dtype=tf.float32)
indexed_positions = tf.cast(
tf.multiply(
tf.cumsum(selected_positions), selected_positions),
dtype=tf.int32)
one_hot_selector = tf.one_hot(
indexed_positions - 1, indicator_sum, dtype=tf.float32)
sampled_indices = tf.cast(
tf.tensordot(
tf.cast(tf.range(tf.shape(indicator)[0]), dtype=tf.float32),
one_hot_selector,
axes=[0, 0]),
dtype=tf.int32)
return gather(boxlist, sampled_indices, use_static_shapes=True)
else:
subboxlist = box_list.BoxList(tf.boolean_mask(boxlist.get(), indicator))
if fields is None:
fields = boxlist.get_extra_fields()
for field in fields:
if not boxlist.has_field(field):
raise ValueError('boxlist must contain all specified fields')
subfieldlist = tf.boolean_mask(boxlist.get_field(field), indicator)
subboxlist.add_field(field, subfieldlist)
return subboxlist
def oscillator_bank(frequency_envelopes: tf.Tensor,
amplitude_envelopes: tf.Tensor,
sample_rate: int = 16000) -> tf.Tensor:
"""Generates audio from sample-wise frequencies for a bank of oscillators.
Args:
frequency_envelopes: Sample-wise oscillator frequencies (Hz). Shape
[batch_size, n_samples, n_sinusoids].
amplitude_envelopes: Sample-wise oscillator amplitude. Shape [batch_size,
n_samples, n_sinusoids].
sample_rate: Sample rate in samples per a second.
Returns:
wav: Sample-wise audio. Shape [batch_size, n_samples, n_sinusoids].
"""
frequency_envelopes = tf_float32(frequency_envelopes)
amplitude_envelopes = tf_float32(amplitude_envelopes)
# Don't exceed Nyquist.
amplitude_envelopes = remove_above_nyquist(frequency_envelopes,
amplitude_envelopes,
sample_rate)
# Change Hz to radians per sample.
omegas = frequency_envelopes * (2.0 * np.pi) # rad / sec
omegas = omegas / float(sample_rate) # rad / sample
# Accumulate phase and synthesize.
phases = tf.cumsum(omegas, axis=1)
wavs = tf.sin(phases)
harmonic_audio = amplitude_envelopes * wavs # [mb, n_samples, n_sinusoids]
audio = tf.reduce_sum(harmonic_audio, axis=-1) # [mb, n_samples]
return audio
def testGradient(self):
x = tf.convert_to_tensor(np.arange(10)[np.newaxis, ...] / 10.0 - 0.5,
dtype=tf.float64)
jac_naive = batch_jacobian(lambda t: tf.cumsum(tf.exp(t), axis=-1), x)
jac_fused = batch_jacobian(
lambda t: tf.exp(tfp.math.log_cumsum_exp(t, axis=-1)), x)
self.assertAllClose(jac_naive, jac_fused)
def _log_prob(self, x):
scores = tf.convert_to_tensor(self.scores)
event_size = self._event_size(scores)
x = tf.cast(x, self.dtype)
# Broadcast scores or x if need be.
if (not tensorshape_util.is_fully_defined(x.shape)
or not tensorshape_util.is_fully_defined(scores.shape)
or x.shape != scores.shape):
broadcast_shape = tf.broadcast_dynamic_shape(
tf.shape(scores), tf.shape(x))
scores = tf.broadcast_to(scores, broadcast_shape)
x = tf.broadcast_to(x, broadcast_shape)
scores_shape = tf.shape(scores)[:-1]
scores_2d = tf.reshape(scores, [-1, event_size])
x_2d = tf.reshape(x, [-1, event_size])
rearranged_scores = tf.gather(scores_2d, x_2d, batch_dims=1)
normalization_terms = tf.cumsum(rearranged_scores,
axis=-1,
reverse=True)
ret = tf.math.reduce_sum(tf.math.log(rearranged_scores /
normalization_terms),
axis=-1)
# Reshape back to user-supplied batch and sample dims prior to 2D reshape.
ret = tf.reshape(ret, scores_shape)
return ret
def cumsum(a, axis=None, dtype=None):
"""Returns cumulative sum of `a` along an axis or the flattened array.
Uses `tf.cumsum`.
Args:
a: array_like. Could be an ndarray, a Tensor or any object that can
be converted to a Tensor using `tf.convert_to_tensor`.
axis: Optional. Axis along which to compute sums. If None, operation is
performed on the flattened array.
dtype: Optional. The type of the output array. If None, defaults to the
dtype of `a` unless `a` is an integer type with precision less than `int`
in which case the output type is `int.`
Returns:
An ndarray with the same number of elements as `a`. If `axis` is None, the
output is a 1-d array, else it has the same shape as `a`.
"""
a = array_creation.asarray(a, dtype=dtype)
if dtype is None and tf.as_dtype(a.dtype).is_integer:
# If a is an integer type and its precision is less than that of `int`,
# the output type will be `int`.
output_type = np.promote_types(a.dtype, int)
if output_type != a.dtype:
a = array_creation.asarray(a, dtype=output_type)
# If axis is None, the input is flattened.
if axis is None:
a = ravel(a)
axis = 0
if axis < 0:
axis += a.ndim
assert axis >= 0 and axis < a.ndim
return utils.tensor_to_ndarray(tf.cumsum(a.data, axis))
def _get_values_from_start_and_end(self, input_tensor, num_start_samples,
num_end_samples, total_num_samples):
"""slices num_start_samples and last num_end_samples from input_tensor.
Args:
input_tensor: An int32 tensor of shape [N] to be sliced.
num_start_samples: Number of examples to be sliced from the beginning
of the input tensor.
num_end_samples: Number of examples to be sliced from the end of the
input tensor.
total_num_samples: Sum of is num_start_samples and num_end_samples. This
should be a scalar.
Returns:
A tensor containing the first num_start_samples and last num_end_samples
from input_tensor.
"""
input_length = tf.shape(input=input_tensor)[0]
start_positions = tf.less(tf.range(input_length), num_start_samples)
end_positions = tf.greater_equal(tf.range(input_length),
input_length - num_end_samples)
selected_positions = tf.logical_or(start_positions, end_positions)
selected_positions = tf.cast(selected_positions, tf.float32)
indexed_positions = tf.multiply(tf.cumsum(selected_positions),
selected_positions)
one_hot_selector = tf.one_hot(tf.cast(indexed_positions, tf.int32) - 1,
total_num_samples,
dtype=tf.float32)
return tf.cast(
tf.tensordot(tf.cast(input_tensor, tf.float32),
one_hot_selector,
axes=[0, 0]), tf.int32)
def quantiles(self):
maxs = self._maxima[:, tf.newaxis]
mins = self._minima[:, tf.newaxis]
weights = tf.cumsum(tf.nn.softmax(self._q, axis=-1), axis=-1)
weights = tf.concat([tf.zeros((self._num_features, 1)), weights],
axis=1)
return weights * (maxs - mins) + mins
def _cdf(self, k):
k = tf.convert_to_tensor(value=k, name="k")
k, probs = _broadcast_cat_event_and_params(
k, self.probs, base_dtype=dtype_util.base_dtype(self.dtype))
# Since the lowest number in the support is 0, any k < 0 should be zero in
# the output.
should_be_zero = k < 0
# Will use k as an index in the gather below, so clip it to {0,...,K-1}.
k = tf.clip_by_value(tf.cast(k, tf.int32), 0, self.num_categories - 1)
batch_shape = tf.shape(input=k)
# tf.gather(..., batch_dims=batch_dims) requires static batch_dims kwarg, so
# to handle the case where the batch shape is dynamic, flatten the batch
# dims (so we know batch_dims=1).
k_flat_batch = tf.reshape(k, [-1])
probs_flat_batch = tf.reshape(
probs, tf.concat(([-1], [self.num_categories]), axis=0))
cdf_flat = tf.gather(tf.cumsum(probs_flat_batch, axis=-1),
k_flat_batch[..., tf.newaxis],
batch_dims=1)
cdf = tf.reshape(cdf_flat, shape=batch_shape)
return tf.where(should_be_zero, tf.zeros_like(cdf), cdf)
def _piecewise_constant_integrate(x1, x2, jump_locations, values, batch_rank):
"""Integrates piecewise constant function between `x1` and `x2`."""
# Initializer already verified that `jump_locations` and `values` have the
# same shape.
# Expand batch size to one if there is no batch shape.
if x1.shape.as_list()[:batch_rank]:
no_batch_shape = False
else:
no_batch_shape = True
x1 = tf.expand_dims(x1, 0)
x2 = tf.expand_dims(x2, 0)
if not jump_locations.shape.as_list()[:-1]:
jump_locations = tf.expand_dims(jump_locations, 0)
values = tf.expand_dims(values, 0)
batch_rank += 1
# Compute integral values between the jump locations
event_shape = tf.shape(values)[(batch_rank + 1):]
event_rank = values.shape.rank - batch_rank - 1
num_data_points = tf.shape(values)[batch_rank]
diff = jump_locations[..., 1:] - jump_locations[..., :-1]
# Broadcast `diff` to the shape of
# `batch_shape + [num_data_points - 2] + [1] * sample_rank`.
for _ in range(event_rank):
diff = tf.expand_dims(diff, -1)
slice_indices = batch_rank * [slice(None)]
slice_indices += [slice(1, num_data_points - 1)]
integrals = tf.cumsum(values[slice_indices] * diff, batch_rank)
# Pad integrals with zero values on left and right.
batch_shape = tf.shape(integrals)[:batch_rank]
pad_shape = tf.concat([batch_shape, [1], event_shape], axis=0)
zeros = tf.zeros(pad_shape, dtype=integrals.dtype)
integrals = tf.concat([zeros, integrals, zeros], axis=batch_rank)
# Get jump locations and values and the integration end points
value1, jump_location1, indices_1 = _get_indices_and_values(
x1, jump_locations, values, 'left', batch_rank)
value2, jump_location2, indices_2 = _get_indices_and_values(
x2, jump_locations, values, 'right', batch_rank)
integrals1 = tf.gather(integrals,
indices_1,
axis=batch_rank,
batch_dims=batch_rank)
integrals2 = tf.gather(integrals,
indices_2,
axis=batch_rank,
batch_dims=batch_rank)
# Broadcast `x1`, `x2`, `jump_location1`, `jump_location2` to the shape
# `batch_shape + [num_points] + [1] * sample_rank`.
for _ in range(event_rank):
x1 = tf.expand_dims(x1, -1)
x2 = tf.expand_dims(x2, -1)
jump_location1 = tf.expand_dims(jump_location1, -1)
jump_location2 = tf.expand_dims(jump_location2, -1)
# Compute the value of the integrals.
res = ((jump_location1 - x1) * value1 + (x2 - jump_location2) * value2 +
integrals2 - integrals1)
if no_batch_shape:
return tf.squeeze(res, 0)
else:
return res
def _joint_sample_n(self, n, seed=None):
"""Draw a joint sample from the prior over latents and observations.
This sampler is specific to LocalLevel models and is faster than the
generic LinearGaussianStateSpaceModel implementation.
Args:
n: `int` `Tensor` number of samples to draw.
seed: Optional `int` `Tensor` seed for the random number generator.
Returns:
latents: `float` `Tensor` of shape `concat([[n], self.batch_shape,
[self.num_timesteps, self.latent_size]], axis=0)` representing samples
of latent trajectories.
observations: `float` `Tensor` of shape `concat([[n], self.batch_shape,
[self.num_timesteps, self.observation_size]], axis=0)` representing
samples of observed series generated from the sampled `latents`.
"""
with tf.name_scope('joint_sample_n'):
(initial_level_seed, level_jumps_seed,
prior_observation_seed) = samplers.split_seed(
seed, n=3, salt='LocalLevelStateSpaceModel_joint_sample_n')
if self.batch_shape.is_fully_defined():
batch_shape = self.batch_shape.as_list()
else:
batch_shape = self.batch_shape_tensor()
sample_and_batch_shape = tf.cast(
prefer_static.concat([[n], batch_shape], axis=0), tf.int32)
# Sample the initial timestep from the prior. Since we want
# this sample to have full batch shape (not just the batch shape
# of the self.initial_state_prior object which might in general be
# smaller), we augment the sample shape to include whatever
# extra batch dimensions are required.
initial_level = self.initial_state_prior.sample(
linear_gaussian_ssm._augment_sample_shape( # pylint: disable=protected-access
self.initial_state_prior, sample_and_batch_shape,
self.validate_args),
seed=initial_level_seed)
# Sample the latent random walk and observed noise, more efficiently than
# the generic loop in `LinearGaussianStateSpaceModel`.
level_jumps = self.level_scale[..., tf.newaxis] * samplers.normal(
prefer_static.concat(
[sample_and_batch_shape, [self.num_timesteps - 1]],
axis=0),
dtype=self.dtype,
seed=level_jumps_seed)
prior_level_sample = tf.cumsum(tf.concat(
[initial_level, level_jumps], axis=-1),
axis=-1)
prior_observation_sample = prior_level_sample + ( # Sample noise.
self.observation_noise_scale[..., tf.newaxis] *
samplers.normal(prefer_static.shape(prior_level_sample),
dtype=self.dtype,
seed=prior_observation_seed))
return (prior_level_sample[..., tf.newaxis],
prior_observation_sample[..., tf.newaxis])
def test_cumulative_sum_power_of_two(self):
elems = self._maybe_static(tf.range(0, 2**4, dtype=tf.int64))
self.assertAllEqual(
self.evaluate(
tfp.math.scan_associative(operator.add,
elems,
max_num_levels=8)),
self.evaluate(tf.cumsum(elems)))
def test_cumulative_sum_size_zero(self):
elems = tf.range(0, dtype=tf.int64)
self.assertAllEqual(
self.evaluate(
tfp.math.scan_associative(operator.add,
elems,
max_num_levels=8)),
self.evaluate(tf.cumsum(elems)))
def testGradient(self):
x = np.arange(10) / 10.0 - 0.5
x = tf.convert_to_tensor(x, dtype=tf.float64)
grad_naive_theoretical, _ = gradient_checker_v2.compute_gradient(
lambda y: tf.cumsum(tf.exp(y)), [x])
grad_fused_theoretical, _ = gradient_checker_v2.compute_gradient(
lambda y: tf.exp(tfp.math.log_cumsum_exp(y)), [x])
self.assertAllClose(grad_fused_theoretical, grad_naive_theoretical)
def test_soft_quantile_normalization(self):
x = tf.constant([1.2, 1.3, 1.5, -4.0, 1.8, 2.4, -1.0])
target = tf.cumsum(tf.ones(x.shape[0]))
xn = ops.soft_quantile_normalization(x, target)
# Make sure that the order of x and xn are identical
self.assertAllEqual(tf.argsort(x), tf.argsort(xn))
# Make sure that the values of xn and target are close.
self.assertAllClose(tf.sort(target), tf.sort(xn), atol=1e-1)
def build_lut(histo, step):
# Compute the cumulative sum, shifting by step // 2
# and then normalization by step.
lut = (tf.cumsum(histo) + (step // 2)) // step
# Shift lut, prepending with 0.
lut = tf.concat([[0], lut[:-1]], 0)
# Clip the counts to be in range. This is done
# in the C code for image.point.
return tf.clip_by_value(lut, 0, 255)
def test_can_scan_tensors_of_different_rank(self):
num_elems = 2**4
elems0 = self.evaluate(
tfd.Uniform(-1., 1.).sample(sample_shape=[num_elems]))
elems1 = self.evaluate(
tfd.Uniform(-1., 1.).sample(sample_shape=[num_elems, 1]))
def extended_add(a, b):
return (a[0] + b[0], a[1] + b[1])
result = self.evaluate(
tfp.math.scan_associative(
extended_add,
(self._maybe_static(elems0), self._maybe_static(elems1))))
self.assertAllClose(result[0], self.evaluate(tf.cumsum(elems0)))
self.assertAllClose(result[1], self.evaluate(tf.cumsum(elems1,
axis=0)))
def _piecewise_constant_integrate(x1, x2, jump_locations, values, batch_rank):
"""Integrates piecewise constant function between `x1` and `x2`."""
# Initializer already verified that `jump_locations` and `values` have the
# same shape.
# Expand batch size to one if there is no batch shape.
if x1.shape.as_list()[:batch_rank]:
no_batch_shape = False
else:
no_batch_shape = True
x1 = tf.expand_dims(x1, 0)
x2 = tf.expand_dims(x2, 0)
if not jump_locations.shape.as_list()[:-1]:
jump_locations = tf.expand_dims(jump_locations, 0)
values = tf.expand_dims(values, 0)
batch_rank += 1
# Compute the index matrix that is later used for `tf.gather_nd`.
index_matrix = _prepare_index_matrix(
x1.shape.as_list()[:-1], x1.shape.as_list()[-1], tf.int32)
# Compute integral values at the jump locations starting from the first jump
# location.
event_shape = values.shape[(batch_rank+1):]
num_data_points = values.shape.as_list()[batch_rank]
diff = jump_locations[..., 1:] - jump_locations[..., :-1]
# Broadcast `diff` to the shape of
# `batch_shape + [num_data_points - 2] + [1] * sample_rank`.
for _ in event_shape:
diff = tf.expand_dims(diff, -1)
slice_indices = batch_rank * [slice(None)]
slice_indices += [slice(1, num_data_points - 1)]
integrals = tf.cumsum(values[slice_indices] * diff, batch_rank)
# Pad integrals with zero values on left and right.
batch_shape = integrals.shape.as_list()[:batch_rank]
zeros = tf.zeros(batch_shape + [1] + event_shape, dtype=integrals.dtype)
integrals = tf.concat([zeros, integrals, zeros], axis=batch_rank)
# Get jump locations and values and the integration end points
value1, jump_location1, indices_nd1 = _get_indices_and_values(
x1, index_matrix, jump_locations, values, 'left', batch_rank)
value2, jump_location2, indices_nd2 = _get_indices_and_values(
x2, index_matrix, jump_locations, values, 'right', batch_rank)
integrals1 = tf.gather_nd(integrals, indices_nd1)
integrals2 = tf.gather_nd(integrals, indices_nd2)
# Broadcast `x1`, `x2`, `jump_location1`, `jump_location2` to the shape
# `batch_shape + [num_points] + [1] * sample_rank`.
for _ in event_shape:
x1 = tf.expand_dims(x1, -1)
x2 = tf.expand_dims(x2, -1)
jump_location1 = tf.expand_dims(jump_location1, -1)
jump_location2 = tf.expand_dims(jump_location2, -1)
# Compute the value of the integral.
res = ((jump_location1 - x1) * value1
+ (x2 - jump_location2) * value2
+ integrals2 - integrals1)
if no_batch_shape:
return tf.squeeze(res, 0)
else:
return res
def test_cumulative_sum_with_xla(self):
elems = self._maybe_static(tf.range(0, 2**4 - 1, dtype=np.int64))
xla_scan = tf.function(experimental_compile=True)(
tfp.math.scan_associative)
result = xla_scan(operator.add, elems)
self.assertAllEqual(self.evaluate(result),
self.evaluate(tf.cumsum(elems)))
def test_max_allowed_size(self):
elems = self.evaluate(tfd.Uniform(-1., 1.).sample([511]))
result = self.evaluate(
tfp.math.scan_associative(operator.add,
self._maybe_static(elems),
max_num_levels=8,
validate_args=True))
self.assertAllClose(result, self.evaluate(tf.cumsum(elems)), atol=1e-4)
def test_cumulative_sum_maximally_odd(self):
# A size that is one less than a power of two ensures that
# every reduction results in an odd size tensor.
# This makes a good test for the logic to handle
# odd sizes
elems = self._maybe_static(tf.range(0, 2**4 - 1, dtype=np.int64))
self.assertAllEqual(
self.evaluate(tfp.math.scan_associative(operator.add, elems)),
self.evaluate(tf.cumsum(elems)))
def _checkBijectorInAllDims(self, x):
"""Helper for `testBijector`."""
x = self._build_tensor(x)
for axis in range(-self.evaluate(tf.rank(x)), 0):
bijector = tfb.Cumsum(axis=axis, validate_args=True)
self.assertStartsWith(bijector.name, 'cumsum')
y = tf.cumsum(x, axis=axis)
self.assertAllClose(y, self.evaluate(bijector.forward(x)))
self.assertAllClose(x, self.evaluate(bijector.inverse(y)))
def _get_new_item_indices(self, age, updates, mask=None):
any_update = list(updates.values())[0]
num_updates = tf.shape(any_update)[0]
_, new_item_indices = tf.math.top_k(age, num_updates)
if mask is not None:
mask = tf.cast(mask, dtype=tf.int32)
unmasked_indices = (tf.cumsum(mask) - 1) * mask
unmasked_indices = tf.expand_dims(unmasked_indices, axis=1)
new_item_indices = tf.gather_nd(new_item_indices, unmasked_indices)
return new_item_indices
def get_note_mask(q_pitch, max_regions=100, note_on_only=True):
"""Get a binary mask for each note from a monophonic instrument.
Each transition of the value creates a new region. Returns the mask of each
region.
Args:
q_pitch: A quantized value, such as pitch or velocity. Shape
[batch, n_timesteps] or [batch, n_timesteps, 1].
max_regions: Maximum number of note regions to consider in the sequence.
Also, the channel dimension of the output mask. Each value transition
defines a new region, e.g. each note-on and note-off count as a separate
region.
note_on_only: Return a mask that is true only for regions where the pitch
is greater than 0.
Returns:
A binary mask of each region [batch, n_timesteps, max_regions].
"""
# Only batch and time dimensions.
if len(q_pitch.shape) == 3:
q_pitch = q_pitch[:, :, 0]
# Get onset and offset points.
edges = tf.abs(spectral_ops.diff(q_pitch, axis=1)) > 0
# Count endpoints as starts/ends of regions.
edges = edges[:, :-1, ...]
edges = tf.pad(edges, [[0, 0], [1, 0]],
mode='constant',
constant_values=True)
edges = tf.pad(edges, [[0, 0], [0, 1]],
mode='constant',
constant_values=False)
edges = tf.cast(edges, tf.int32)
# Count up onset and offsets for each timestep.
# Assumes each onset has a corresponding offset.
# The -1 ensures that the 0th index is the first note.
edge_idx = tf.cumsum(edges, axis=1) - 1
# Create masks of shape [batch, n_timesteps, max_regions].
note_mask = edge_idx[..., None] == tf.range(max_regions)[None, None, :]
note_mask = tf.cast(note_mask, tf.float32)
if note_on_only:
# [batch, notes]
note_pitches = get_note_moments(q_pitch, note_mask, return_std=False)
# [batch, time, notes]
note_on = tf.cast(note_pitches > 0.0, tf.float32)[:, None, :]
# [batch, time, notes]
note_mask *= note_on
return note_mask
def cumsum(a, axis=None, dtype=None): # pylint: disable=missing-docstring
a = asarray(a, dtype=dtype)
if dtype is None:
a = _maybe_promote_to_int(a)
# If axis is None, the input is flattened.
if axis is None:
a = ravel(a)
axis = 0
elif axis < 0:
axis += tf.rank(a.data)
return utils.tensor_to_ndarray(tf.cumsum(a.data, axis))
def test_cumulative_sum_custom_axis(self, axis):
elems = self._maybe_static(
tf.random.stateless_normal(
[4, 32, 31, 1],
seed=test_util.test_seed(sampler_type='stateless')))
axis = self._maybe_static(axis)
expected_result = self.evaluate(tf.cumsum(elems, axis=axis))
result = tfp.math.scan_associative(operator.add,
elems,
axis=axis,
max_num_levels=8)
self.assertAllClose(self.evaluate(result), expected_result, rtol=1e-5)
