def crf_unary_score(tag_indices, sequence_lengths, inputs):
"""Computes the unary scores of tag sequences.
Args:
tag_indices: A [batch_size, max_seq_len] matrix of tag indices.
sequence_lengths: A [batch_size] vector of true sequence lengths.
inputs: A [batch_size, max_seq_len, num_tags] tensor of unary potentials.
Returns:
unary_scores: A [batch_size] vector of unary scores.
"""
batch_size = array_ops.shape(inputs)[0]
max_seq_len = array_ops.shape(inputs)[1]
num_tags = array_ops.shape(inputs)[2]
flattened_inputs = array_ops.reshape(inputs, [-1])
offsets = array_ops.expand_dims(
math_ops.range(batch_size) * max_seq_len * num_tags, 1)
offsets += array_ops.expand_dims(math_ops.range(max_seq_len) * num_tags, 0)
flattened_tag_indices = array_ops.reshape(offsets + tag_indices, [-1])
unary_scores = array_ops.reshape(
array_ops.gather(flattened_inputs, flattened_tag_indices),
[batch_size, max_seq_len])
masks = _lengths_to_masks(sequence_lengths, array_ops.shape(tag_indices)[1])
unary_scores = math_ops.reduce_sum(unary_scores * masks, 1)
return unary_scores
def _fn():
num_rows = np.shape(np_matrix)[0]
num_cols = np.shape(np_matrix)[1]
row_ids = math_ops.range(num_rows, dtype=dtypes.int64)
col_ids = math_ops.range(num_cols, dtype=dtypes.int64)
sp_mat = self.np_array_to_sparse(np_matrix)
sp_mat_t = sparse_ops.sparse_transpose(sp_mat)
row_batch = input_lib.batch(
[row_ids, sp_mat],
batch_size=min(batch_size, num_rows),
capacity=10,
enqueue_many=True)
col_batch = input_lib.batch(
[col_ids, sp_mat_t],
batch_size=min(batch_size, num_cols),
capacity=10,
enqueue_many=True)
features = extract_features(row_batch, col_batch, sp_mat.dense_shape)
if projection_weights is not None:
weights_batch = input_lib.batch(
projection_weights,
batch_size=batch_size,
capacity=10,
enqueue_many=True)
features[wals_lib.WALSMatrixFactorization.PROJECTION_WEIGHTS] = (
weights_batch)
if project_row is not None:
features[wals_lib.WALSMatrixFactorization.PROJECT_ROW] = (
constant_op.constant(project_row))
labels = None
return features, labels
def frames(signal, frame_length, frame_step, name=None):
"""Frame a signal into overlapping frames.
May be used in front of spectral functions.
For example:
```python
pcm = tf.placeholder(tf.float32, [None, 9152])
frames = tf.contrib.signal.frames(pcm, 512, 180)
magspec = tf.abs(tf.spectral.rfft(frames, [512]))
image = tf.expand_dims(magspec, 3)
```
Args:
signal: A `Tensor` of shape `[batch_size, signal_length]`.
frame_length: An `int32` or `int64` `Tensor`. The length of each frame.
frame_step: An `int32` or `int64` `Tensor`. The step between frames.
name: A name for the operation (optional).
Returns:
A `Tensor` of frames with shape `[batch_size, num_frames, frame_length]`.
Raises:
ValueError: if signal does not have rank 2.
"""
with ops.name_scope(name, "frames", [signal, frame_length, frame_step]):
signal = ops.convert_to_tensor(signal, name="signal")
frame_length = ops.convert_to_tensor(frame_length, name="frame_length")
frame_step = ops.convert_to_tensor(frame_step, name="frame_step")
signal_rank = signal.shape.ndims
if signal_rank != 2:
raise ValueError("expected signal to have rank 2 but was " + signal_rank)
signal_length = array_ops.shape(signal)[1]
num_frames = math_ops.ceil((signal_length - frame_length) / frame_step)
num_frames = 1 + math_ops.cast(num_frames, dtypes.int32)
pad_length = (num_frames - 1) * frame_step + frame_length
pad_signal = array_ops.pad(signal, [[0, 0], [0,
pad_length - signal_length]])
indices_frame = array_ops.expand_dims(math_ops.range(frame_length), 0)
indices_frames = array_ops.tile(indices_frame, [num_frames, 1])
indices_step = array_ops.expand_dims(
math_ops.range(num_frames) * frame_step, 1)
indices_steps = array_ops.tile(indices_step, [1, frame_length])
indices = indices_frames + indices_steps
# TODO(androbin): remove `transpose` when `gather` gets `axis` support
pad_signal = array_ops.transpose(pad_signal)
signal_frames = array_ops.gather(pad_signal, indices)
signal_frames = array_ops.transpose(signal_frames, perm=[2, 0, 1])
return signal_frames
def test_docstring_example(self):
# Produce the first 1000 members of the Halton sequence in 3 dimensions.
num_results = 1000
dim = 3
with self.test_session():
sample = halton.sample(dim, num_results=num_results, randomized=False)
# Evaluate the integral of x_1 * x_2^2 * x_3^3 over the three dimensional
# hypercube.
powers = math_ops.range(1.0, limit=dim + 1)
integral = math_ops.reduce_mean(
math_ops.reduce_prod(sample ** powers, axis=-1))
true_value = 1.0 / math_ops.reduce_prod(powers + 1.0)
# Produces a relative absolute error of 1.7%.
self.assertAllClose(integral.eval(), true_value.eval(), rtol=0.02)
# Now skip the first 1000 samples and recompute the integral with the next
# thousand samples. The sequence_indices argument can be used to do this.
sequence_indices = math_ops.range(start=1000, limit=1000 + num_results,
dtype=dtypes.int32)
sample_leaped = halton.sample(dim, sequence_indices=sequence_indices,
randomized=False)
integral_leaped = math_ops.reduce_mean(
math_ops.reduce_prod(sample_leaped ** powers, axis=-1))
self.assertAllClose(integral_leaped.eval(), true_value.eval(), rtol=0.05)
def _tf_range(start_or_stop, stop, step):
# TODO(mdan): We should optimize this when a full tensor is not required.
if step is not UNDEFINED:
return math_ops.range(start_or_stop, stop, step)
if stop is not UNDEFINED:
return math_ops.range(start_or_stop, stop)
return math_ops.range(start_or_stop)
def move_right_permutation():
return util.prefer_static_value(
array_ops.concat([
math_ops.range(0, source_idx, dtype=dtype),
math_ops.range(source_idx+1, dest_idx+1, dtype=dtype),
[source_idx],
math_ops.range(dest_idx+1, ndims, dtype=dtype)], axis=0))
def _do_maximum_mean(samples, envelope, high, name=None):
"""Common code between maximum_mean and minimum_mean."""
with ops.name_scope(name, "do_maximum_mean", [samples, envelope, high]):
n = array_ops.rank(samples)
# Move the batch dimension of `samples` to the rightmost position,
# where the _batch_sort_vector function wants it.
perm = array_ops.concat([math_ops.range(1, n), [0]], axis=0)
samples = array_ops.transpose(samples, perm)
samples = _batch_sort_vector(samples)
# The maximum mean is given by taking `envelope`-worth of
# probability from the smallest samples and moving it to the
# maximum value. This amounts to:
# - ignoring the smallest k samples, where `k/n < envelope`
# - taking a `1/n - (envelope - k/n)` part of the index k sample
# - taking all the other samples
# - and adding `envelope * high` at the end.
# The following is a vectorized and batched way of computing this.
# `max_mean_contrib` is a mask implementing the previous.
batch_size = array_ops.shape(samples)[-1]
batch_size = math_ops.cast(batch_size, dtype=samples.dtype.base_dtype)
step = 1. / batch_size
cum_steps = step * math_ops.range(
1, batch_size + 1, dtype=samples.dtype.base_dtype)
max_mean_contrib = clip_ops.clip_by_value(
cum_steps - envelope[..., array_ops.newaxis],
clip_value_min=0.,
clip_value_max=step)
return math_ops.reduce_sum(
samples * max_mean_contrib, axis=-1) + envelope * high
def _sample_n(self, n, seed):
batch_shape = self.batch_shape_tensor()
event_shape = self.event_shape_tensor()
batch_ndims = array_ops.shape(batch_shape)[0]
ndims = batch_ndims + 3 # sample_ndims=1, event_ndims=2
shape = array_ops.concat([[n], batch_shape, event_shape], 0)
# Complexity: O(nbk**2)
x = random_ops.random_normal(shape=shape,
mean=0.,
stddev=1.,
dtype=self.dtype,
seed=seed)
# Complexity: O(nbk)
# This parametrization is equivalent to Chi2, i.e.,
# ChiSquared(k) == Gamma(alpha=k/2, beta=1/2)
expanded_df = self.df * array_ops.ones(
self.scale_operator.batch_shape_tensor(),
dtype=self.df.dtype.base_dtype)
g = random_ops.random_gamma(shape=[n],
alpha=self._multi_gamma_sequence(
0.5 * expanded_df, self.dimension),
beta=0.5,
dtype=self.dtype,
seed=distribution_util.gen_new_seed(
seed, "wishart"))
# Complexity: O(nbk**2)
x = array_ops.matrix_band_part(x, -1, 0) # Tri-lower.
# Complexity: O(nbk)
x = array_ops.matrix_set_diag(x, math_ops.sqrt(g))
# Make batch-op ready.
# Complexity: O(nbk**2)
perm = array_ops.concat([math_ops.range(1, ndims), [0]], 0)
x = array_ops.transpose(x, perm)
shape = array_ops.concat([batch_shape, [event_shape[0]], [-1]], 0)
x = array_ops.reshape(x, shape)
# Complexity: O(nbM) where M is the complexity of the operator solving a
# vector system. E.g., for LinearOperatorDiag, each matmul is O(k**2), so
# this complexity is O(nbk**2). For LinearOperatorLowerTriangular,
# each matmul is O(k^3) so this step has complexity O(nbk^3).
x = self.scale_operator.matmul(x)
# Undo make batch-op ready.
# Complexity: O(nbk**2)
shape = array_ops.concat([batch_shape, event_shape, [n]], 0)
x = array_ops.reshape(x, shape)
perm = array_ops.concat([[ndims - 1], math_ops.range(0, ndims - 1)], 0)
x = array_ops.transpose(x, perm)
if not self.cholesky_input_output_matrices:
# Complexity: O(nbk^3)
x = math_ops.matmul(x, x, adjoint_b=True)
return x
def to_weighted_sum(self,
input_tensor,
num_outputs=1,
weight_collections=None,
trainable=True):
"""Returns a Tensor as linear predictions and a list of created Variable."""
dimension = self.source_column.dimension
batch_size = array_ops.shape(input_tensor)[0]
if dimension > 1:
i1 = array_ops.reshape(array_ops.tile(array_ops.expand_dims(
math_ops.range(0, batch_size), 1), [1, dimension]), [-1])
i2 = array_ops.tile(math_ops.range(0, dimension), [batch_size])
# Flatten the bucket indices and unique them across dimensions
# E.g. 2nd dimension indices will range from k to 2*k-1 with k buckets
# TODO(chapelle): move that logic to insert_transformed_feature to ensure
# unique buckets across dimensions after crossing.
bucket_indices = array_ops.reshape(input_tensor, [-1]) + self.length * i2
else:
# Simpler indices when dimension=1
i1 = math_ops.range(0, batch_size)
i2 = array_ops.zeros([batch_size], dtype=dtypes.int32)
bucket_indices = array_ops.reshape(input_tensor, [-1])
indices = math_ops.to_int64(array_ops.transpose(array_ops.pack((i1, i2))))
shape = math_ops.to_int64(array_ops.pack([batch_size, 1]))
sparse_id_values = ops.SparseTensor(indices, bucket_indices, shape)
vocab_size = self.length * self.source_column.dimension
return _create_embedding_lookup(
sparse_id_values, vocab_size, num_outputs,
_add_variable_collection(weight_collections), 0., "sum",
trainable, self.name + "_weights")
def _ExtractImagePatchesGrad(op, grad):
batch_size, rows_in, cols_in, channels = [
dim.value for dim in op.inputs[0].shape.dims
]
input_bhwc = array_ops.shape(op.inputs[0])
batch_size = input_bhwc[0]
channels = input_bhwc[3]
# Create indices matrix for input tensor.
# Note that 0 is preserved for padding location,
# so indices for input start from 1 to 1 + rows_in * cols_in.
input_indices_num = 1 + rows_in * cols_in
input_idx = array_ops.reshape(math_ops.range(1, input_indices_num,
dtype=ops.dtypes.int64),
(1, rows_in, cols_in, 1))
input_idx_patched = gen_array_ops.extract_image_patches(
input_idx,
op.get_attr("ksizes"),
op.get_attr("strides"),
op.get_attr("rates"),
op.get_attr("padding"))
# Create indices matrix for output tensor.
_, rows_out, cols_out, _ = [dim.value for dim in op.outputs[0].shape.dims]
_, ksize_r, ksize_c, _ = op.get_attr("ksizes")
# Indices for output start from 0.
output_indices_num = rows_out * cols_out * ksize_r * ksize_c
output_idx = array_ops.reshape(math_ops.range(output_indices_num,
dtype=ops.dtypes.int64),
(1, rows_out, cols_out, ksize_r * ksize_c))
# Construct mapping table for indices: (input -> output).
idx_matrix = array_ops.concat(
[array_ops.expand_dims(input_idx_patched, axis=-1),
array_ops.expand_dims(output_idx, axis=-1)],
axis=-1)
idx_map = array_ops.reshape(idx_matrix, (-1, 2))
sp_shape = (input_indices_num, output_indices_num)
sp_mat_full = sparse_tensor.SparseTensor(
idx_map,
array_ops.ones([output_indices_num], dtype=grad.dtype),
sp_shape)
# Remove all padding locations [0, :].
sp_mat = sparse_ops.sparse_slice(sp_mat_full,
(1, 0),
(input_indices_num - 1, output_indices_num))
grad_expanded = array_ops.transpose(
array_ops.reshape(
grad, (batch_size, rows_out, cols_out, ksize_r, ksize_c, channels)),
(1, 2, 3, 4, 0, 5))
grad_flat = array_ops.reshape(grad_expanded, (-1, batch_size * channels))
jac = sparse_ops.sparse_tensor_dense_matmul(sp_mat, grad_flat)
grad_out = array_ops.reshape(jac, (rows_in, cols_in, batch_size, channels))
grad_out = array_ops.transpose(grad_out, (2, 0, 1, 3))
return [grad_out]
def testFormatOneTensorOneDimVarySummarize(self):
with self.test_session():
tensor = math_ops.range(6)
format_output = string_ops.string_format("{}", tensor, summarize=-1)
out = self.evaluate(format_output)
expected = "[0 1 2 3 4 5]"
self.assertEqual(compat.as_text(out), expected)
with self.test_session():
tensor = math_ops.range(6)
format_output = string_ops.string_format("{}", tensor, summarize=1)
out = self.evaluate(format_output)
expected = "[0 ... 5]"
self.assertEqual(compat.as_text(out), expected)
with self.test_session():
tensor = math_ops.range(6)
format_output = string_ops.string_format("{}", tensor, summarize=2)
out = self.evaluate(format_output)
expected = "[0 1 ... 4 5]"
self.assertEqual(compat.as_text(out), expected)
with self.test_session():
tensor = math_ops.range(6)
format_output = string_ops.string_format("{}", tensor, summarize=10)
out = self.evaluate(format_output)
expected = "[0 1 2 3 4 5]"
self.assertEqual(compat.as_text(out), expected)
def testShapePassedToGradient(self):
with ops.Graph().as_default():
@custom_gradient.custom_gradient
def differentiable_scatter_update(handle, indices, values):
with ops.control_dependencies([
resource_variable_ops.resource_scatter_update(
handle, indices, values)]):
new_handle = array_ops.identity(handle)
def grad(dresult):
self.assertIsNotNone(
tensor_util.constant_value(dresult.dense_shape))
return [dresult, None, None]
return new_handle, grad
var = variable_scope.get_variable(
"foo", shape=[20], initializer=init_ops.zeros_initializer,
dtype=dtypes.float64, use_resource=True)
indices = math_ops.range(10)
updates = math_ops.range(9, -1, -1, dtype=dtypes.float64)
new_handle = differentiable_scatter_update(var.handle, indices, updates)
gathered = resource_variable_ops.resource_gather(
new_handle, indices, dtype=var.dtype)
gradients_impl.gradients([gathered], [updates])
def _potential_scale_reduction_single_state(state, independent_chain_ndims):
"""potential_scale_reduction for one single state `Tensor`."""
# We assume exactly one leading dimension indexes e.g. correlated samples from
# each Markov chain.
state = ops.convert_to_tensor(state, name="state")
sample_ndims = 1
sample_axis = math_ops.range(0, sample_ndims)
chain_axis = math_ops.range(sample_ndims,
sample_ndims + independent_chain_ndims)
sample_and_chain_axis = math_ops.range(0,
sample_ndims + independent_chain_ndims)
n = _axis_size(state, sample_axis)
m = _axis_size(state, chain_axis)
# In the language of [2],
# 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(
math_ops.reduce_mean(state, sample_axis, keepdims=True),
sample_and_chain_axis,
biased=False)
w = math_ops.reduce_mean(
_reduce_variance(state, sample_axis, keepdims=True, biased=True),
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 gather_tree_from_array(t, parent_ids, sequence_length):
"""Calculates the full beams for `TensorArray`s.
Args:
t: A stacked `TensorArray` of size `max_time` that contains `Tensor`s of
shape `[batch_size, beam_width, s]` or `[batch_size * beam_width, s]`
where `s` is the depth shape.
parent_ids: The parent ids of shape `[max_time, batch_size, beam_width]`.
sequence_length: The sequence length of shape `[batch_size, beam_width]`.
Returns:
A `Tensor` which is a stacked `TensorArray` of the same size and type as
`t` and where beams are sorted in each `Tensor` according to `parent_ids`.
"""
max_time = parent_ids.shape[0].value or array_ops.shape(parent_ids)[0]
batch_size = parent_ids.shape[1].value or array_ops.shape(parent_ids)[1]
beam_width = parent_ids.shape[2].value or array_ops.shape(parent_ids)[2]
# Generate beam ids that will be reordered by gather_tree.
beam_ids = array_ops.expand_dims(
array_ops.expand_dims(math_ops.range(beam_width), 0), 0)
beam_ids = array_ops.tile(beam_ids, [max_time, batch_size, 1])
mask = array_ops.sequence_mask(
sequence_length, maxlen=max_time, dtype=dtypes.int32)
mask = array_ops.transpose(mask, perm=[2, 0, 1])
# Use beam_width + 1 to mark the end of beam.
masked_beam_ids = (beam_ids * mask) + (1 - mask) * (beam_width + 1)
max_sequence_lengths = math_ops.to_int32(
math_ops.reduce_max(sequence_length, axis=1))
sorted_beam_ids = beam_search_ops.gather_tree(
step_ids=masked_beam_ids,
parent_ids=parent_ids,
max_sequence_lengths=max_sequence_lengths,
end_token=beam_width + 1)
# For out of range steps, simply copy the same beam.
sorted_beam_ids = array_ops.where(
math_ops.cast(mask, dtypes.bool), x=sorted_beam_ids, y=beam_ids)
# Generate indices for gather_nd.
time_ind = array_ops.tile(array_ops.reshape(
math_ops.range(max_time), [-1, 1, 1]), [1, batch_size, beam_width])
batch_ind = array_ops.tile(array_ops.reshape(
math_ops.range(batch_size), [-1, 1, 1]), [1, max_time, beam_width])
batch_ind = array_ops.transpose(batch_ind, perm=[1, 0, 2])
indices = array_ops.stack([time_ind, batch_ind, sorted_beam_ids], -1)
# Gather from a tensor with collapsed additional dimensions.
gather_from = t
final_shape = array_ops.shape(gather_from)
gather_from = array_ops.reshape(
gather_from, [max_time, batch_size, beam_width, -1])
ordered = array_ops.gather_nd(gather_from, indices)
ordered = array_ops.reshape(ordered, final_shape)
return ordered
def _maybe_rotate_dims(self, x, rotate_right=False):
"""Helper which rolls left event_dims left or right event_dims right."""
if tensor_util.constant_value(self._needs_rotation) is False:
return x
ndims = array_ops.rank(x)
n = (ndims - self._rotate_ndims) if rotate_right else self._rotate_ndims
return array_ops.transpose(
x, _concat_vectors(math_ops.range(n, ndims), math_ops.range(0, n)))
def _GatherV2Grad(op, grad):
"""Gradient for GatherV2 op."""
# params can be large, so colocate the shape calculation with it.
#
# params can be very large for sparse model, array_ops.shape raises
# exception on the Windows platform when any dimension is larger than
# int32. params_shape is not used in optimizer apply_sparse gradients,
# so it's fine to convert it back to int32 regardless of truncation.
params = op.inputs[0]
with ops.colocate_with(params):
params_shape = array_ops.shape(params, out_type=ops.dtypes.int64)
params_shape = math_ops.to_int32(params_shape)
indices = op.inputs[1]
indices_size = array_ops.expand_dims(array_ops.size(indices), 0)
axis = op.inputs[2]
axis_static = tensor_util.constant_value(axis)
# For axis 0 gathers, build an appropriately shaped IndexedSlices.
if axis_static == 0:
values_shape = array_ops.concat([indices_size, params_shape[1:]], 0)
values = array_ops.reshape(grad, values_shape)
indices = array_ops.reshape(indices, indices_size)
return [ops.IndexedSlices(values, indices, params_shape), None, None]
outer_shape = params_shape[:axis]
outer_dims = array_ops.size(outer_shape)
inner_shape = params_shape[axis:][1:]
inner_dims = array_ops.size(inner_shape)
outer_axes_indices = math_ops.range(outer_dims)
inner_axes_indices = math_ops.range(outer_dims + 1,
outer_dims + 1 + inner_dims)
values_shape = array_ops.concat([outer_shape, indices_size, inner_shape], 0)
values = array_ops.reshape(grad, values_shape)
indices = array_ops.reshape(indices, indices_size)
# We need to sum up every slice `values[..., i, ....]` corresponding to
# `params[..., indices[i], ...]`. Since `unsorted_segment_sum` does not
# support an axis parameter, we transpose the gather dimension to the front,
# then use `unsorted_segment_sum` to build a
# [gather_axis, outer_axes, inner_axes] tensor with all the gradients
# affecting each index in `gather_axis` summed up.
transpose_dims = array_ops.concat(
[[outer_dims], outer_axes_indices, inner_axes_indices], 0)
values_transpose = array_ops.transpose(values, transpose_dims)
num_segments = params_shape[axis]
params_grad = math_ops.unsorted_segment_sum(
values_transpose, indices, num_segments)
# Inverts the above transpose by moving dimension 0 back to its original
# position.
invert_transpose_dims = array_ops.concat(
[outer_axes_indices + 1, [0], inner_axes_indices], 0)
params_grad = array_ops.transpose(params_grad, invert_transpose_dims)
return [params_grad, None, None]
def _sample_n(self, n, seed):
batch_shape = self.batch_shape()
event_shape = self.event_shape()
batch_ndims = array_ops.shape(batch_shape)[0]
ndims = batch_ndims + 3 # sample_ndims=1, event_ndims=2
shape = array_ops.concat(((n,), batch_shape, event_shape), 0)
# Complexity: O(nbk^2)
x = random_ops.random_normal(shape=shape,
mean=0.,
stddev=1.,
dtype=self.dtype,
seed=seed)
# Complexity: O(nbk)
# This parametrization is equivalent to Chi2, i.e.,
# ChiSquared(k) == Gamma(alpha=k/2, beta=1/2)
g = random_ops.random_gamma(shape=(n,),
alpha=self._multi_gamma_sequence(
0.5 * self.df, self.dimension),
beta=0.5,
dtype=self.dtype,
seed=distribution_util.gen_new_seed(
seed, "wishart"))
# Complexity: O(nbk^2)
x = array_ops.matrix_band_part(x, -1, 0) # Tri-lower.
# Complexity: O(nbk)
x = array_ops.matrix_set_diag(x, math_ops.sqrt(g))
# Make batch-op ready.
# Complexity: O(nbk^2)
perm = array_ops.concat((math_ops.range(1, ndims), (0,)), 0)
x = array_ops.transpose(x, perm)
shape = array_ops.concat((batch_shape, (event_shape[0], -1)), 0)
x = array_ops.reshape(x, shape)
# Complexity: O(nbM) where M is the complexity of the operator solving a
# vector system. E.g., for OperatorPDDiag, each matmul is O(k^2), so
# this complexity is O(nbk^2). For OperatorPDCholesky, each matmul is
# O(k^3) so this step has complexity O(nbk^3).
x = self.scale_operator_pd.sqrt_matmul(x)
# Undo make batch-op ready.
# Complexity: O(nbk^2)
shape = array_ops.concat((batch_shape, event_shape, (n,)), 0)
x = array_ops.reshape(x, shape)
perm = array_ops.concat(((ndims - 1,), math_ops.range(0, ndims - 1)), 0)
x = array_ops.transpose(x, perm)
if not self.cholesky_input_output_matrices:
# Complexity: O(nbk^3)
x = math_ops.matmul(x, x, adjoint_b=True)
return x
def _rotate_last_dim(x, rotate_right=False):
"""Rotate the last dimension either left or right."""
ndims = array_ops.rank(x)
if rotate_right:
transpose_perm = array_ops.concat(
[[ndims - 1], math_ops.range(0, ndims - 1)], axis=0)
else:
transpose_perm = array_ops.concat(
[math_ops.range(1, ndims), [0]], axis=0)
return array_ops.transpose(x, transpose_perm)
def _maybe_rotate_dims(self, x, rotate_right=False):
"""Helper which rolls left event_dims left or right event_dims right."""
if tensor_util.constant_value(self._needs_rotation) is False:
return x
ndims = array_ops.rank(x)
n = _pick_scalar_condition(self._needs_rotation,
self._override_event_ndims, 0)
if rotate_right:
n = ndims - n
return array_ops.transpose(
x, _concat_vectors(math_ops.range(n, ndims), math_ops.range(0, n)))
def _ctc_state_trans(label_seq):
"""Compute CTC alignment model transition matrix.
Args:
label_seq: tensor of shape [batch_size, max_seq_length]
Returns:
tensor of shape [batch_size, states, states] with a state transition matrix
computed for each sequence of the batch.
"""
with ops.name_scope("ctc_state_trans"):
label_seq = ops.convert_to_tensor(label_seq, name="label_seq")
batch_size = _get_dim(label_seq, 0)
num_labels = _get_dim(label_seq, 1)
num_label_states = num_labels + 1
num_states = 2 * num_label_states
label_states = math_ops.range(num_label_states)
blank_states = label_states + num_label_states
# Start state to first label.
start_to_label = [[1, 0]]
# Blank to label transitions.
blank_to_label = array_ops.stack([label_states[1:], blank_states[:-1]], 1)
# Label to blank transitions.
label_to_blank = array_ops.stack([blank_states, label_states], 1)
# Scatter transitions that don't depend on sequence.
indices = array_ops.concat(
[start_to_label, blank_to_label, label_to_blank], 0)
values = array_ops.ones([_get_dim(indices, 0)])
trans = array_ops.scatter_nd(
indices, values, shape=[num_states, num_states])
trans += linalg_ops.eye(num_states) # Self-loops.
# Label to label transitions. Disallow transitions between repeated labels
# with no blank state in between.
batch_idx = array_ops.zeros_like(label_states[2:])
indices = array_ops.stack(
[batch_idx, label_states[2:], label_states[1:-1]], 1)
indices = array_ops.tile(
array_ops.expand_dims(indices, 0), [batch_size, 1, 1])
batch_idx = array_ops.expand_dims(math_ops.range(batch_size), 1) * [1, 0, 0]
indices += array_ops.expand_dims(batch_idx, 1)
repeats = math_ops.equal(label_seq[:, :-1], label_seq[:, 1:])
values = 1.0 - math_ops.cast(repeats, dtypes.float32)
batched_shape = [batch_size, num_states, num_states]
label_to_label = array_ops.scatter_nd(indices, values, batched_shape)
return array_ops.expand_dims(trans, 0) + label_to_label
def dense_image_warp(image, flow, name='dense_image_warp'):
"""Image warping using per-pixel flow vectors.
Apply a non-linear warp to the image, where the warp is specified by a dense
flow field of offset vectors that define the correspondences of pixel values
in the output image back to locations in the source image. Specifically, the
pixel value at output[b, j, i, c] is
images[b, j - flow[b, j, i, 0], i - flow[b, j, i, 1], c].
The locations specified by this formula do not necessarily map to an int
index. Therefore, the pixel value is obtained by bilinear
interpolation of the 4 nearest pixels around
(b, j - flow[b, j, i, 0], i - flow[b, j, i, 1]). For locations outside
of the image, we use the nearest pixel values at the image boundary.
Args:
image: 4-D float `Tensor` with shape `[batch, height, width, channels]`.
flow: A 4-D float `Tensor` with shape `[batch, height, width, 2]`.
name: A name for the operation (optional).
Note that image and flow can be of type tf.half, tf.float32, or tf.float64,
and do not necessarily have to be the same type.
Returns:
A 4-D float `Tensor` with shape`[batch, height, width, channels]`
and same type as input image.
Raises:
ValueError: if height < 2 or width < 2 or the inputs have the wrong number
of dimensions.
"""
with ops.name_scope(name):
batch_size, height, width, channels = (array_ops.shape(image)[0],
array_ops.shape(image)[1],
array_ops.shape(image)[2],
array_ops.shape(image)[3])
# The flow is defined on the image grid. Turn the flow into a list of query
# points in the grid space.
grid_x, grid_y = array_ops.meshgrid(
math_ops.range(width), math_ops.range(height))
stacked_grid = math_ops.cast(
array_ops.stack([grid_y, grid_x], axis=2), flow.dtype)
batched_grid = array_ops.expand_dims(stacked_grid, axis=0)
query_points_on_grid = batched_grid - flow
query_points_flattened = array_ops.reshape(query_points_on_grid,
[batch_size, height * width, 2])
# Compute values at the query points, then reshape the result back to the
# image grid.
interpolated = _interpolate_bilinear(image, query_points_flattened)
interpolated = array_ops.reshape(interpolated,
[batch_size, height, width, channels])
return interpolated
def test_identity(self):
axis_order = ['w', 'x', 'y', 'z']
lt = core.LabeledTensor(
array_ops.reshape(math_ops.range(24), (1, 2, 3, 4)), axis_order)
actual = core.impose_axis_order(lt, axis_order)
self.assertLabeledTensorsEqual(lt, actual)
lt = core.LabeledTensor(
array_ops.reshape(math_ops.range(6), (1, 2, 3)), axis_order[:3])
actual = core.impose_axis_order(lt, axis_order)
self.assertLabeledTensorsEqual(lt, actual)
def _broadcast_x_higher_rank_than_sigma():
x_shape_left = array_ops.slice(x_shape, [0], sigma_rank_vec - 1)
x_shape_right = array_ops.slice(x_shape, sigma_rank_vec - 1, x_rank_vec - 1)
x_shape_perm = array_ops.concat(
0, (math_ops.range(sigma_rank - 1, x_rank), math_ops.range(0, sigma_rank - 1))
)
return array_ops.reshape(
# Convert to [D, E, F, ..., k, B, C]
array_ops.transpose(x_centered, perm=x_shape_perm),
# Reshape to [D, E, F, ..., k, B*C]
array_ops.concat(0, (x_shape_right, array_ops.pack([math_ops.reduce_prod(x_shape_left, 0)]))),
)
def test_matrix_vector(self):
xy_lt = core.LabeledTensor(
array_ops.reshape(math_ops.range(6), (2, 3)), ['x', 'y'])
y_lt = core.LabeledTensor(math_ops.range(3), ['y'])
matmul_lt = ops.matmul(xy_lt, y_lt)
golden_lt = core.LabeledTensor(
math_ops.matmul(xy_lt.tensor, array_ops.reshape(y_lt.tensor,
(-1, 1)))[:, 0], ['x'])
self.assertLabeledTensorsEqual(matmul_lt, golden_lt)
matmul_lt = ops.matmul(y_lt, xy_lt)
self.assertLabeledTensorsEqual(matmul_lt, golden_lt)
def make_tril_ids(n):
"""Internal helper to create vector of linear indices into y."""
cols = array_ops.reshape(array_ops.tile(math_ops.range(n), [n]), [n, n])
rows = array_ops.tile(
array_ops.expand_dims(math_ops.range(n), -1), [1, n])
pred = math_ops.greater(cols, rows)
tril_ids = array_ops.tile(array_ops.reshape(
math_ops.cumsum(math_ops.range(n)), [n, 1]), [1, n]) + cols
tril_ids = math_ops.select(pred,
array_ops.zeros([n, n], dtype=dtypes.int32),
tril_ids + 1)
tril_ids = array_ops.reshape(tril_ids, [-1])
return tril_ids
def _add_sinusoids_signal(x, time, min_timescale=1.0, max_timescale=1.0e4):
"""Adds a bunch of sinusoids of different frequencies to a Tensor.
Each channel of the input Tensor is incremented by a sinusoid of a different
frequency and phase.
This allows attention to learn to use absolute and relative positions.
Timing signals should be added to some precursors of both the query and the
memory inputs to attention.
The use of relative position is possible because sin(x+y) and cos(x+y) can be
experessed in terms of y, sin(x) and cos(x).
In particular, we use a geometric sequence of timescales starting with
min_timescale and ending with max_timescale. The number of different
timescales is equal to channels / 2. For each timescale, we
generate the two sinusoidal signals sin(timestep/timescale) and
cos(timestep/timescale). All of these sinusoids are concatenated in
the channels dimension.
Args:
x: a Tensor with shape [batch, length, channels]
min_timescale: a float
max_timescale: a float
Returns:
a Tensor the same shape as x.
"""
channels = x.get_shape().as_list()[-1]
if x.get_shape().ndims == 3: # [batch_size, timesteps, dim]
length = array_ops.shape(x)[1]
position = math_ops.to_float(math_ops.range(length))
elif x.get_shape().ndims == 2: # [batch_size, dim]
length = 1
position = math_ops.to_float(math_ops.range(time, time + 1))
else:
raise ValueError("need a Tensor with rank 2 or 3")
num_timescales = channels // 2
log_timescale_increment = (
math.log(float(max_timescale) / float(min_timescale)) /
(math_ops.to_float(num_timescales) - 1))
inv_timescales = min_timescale * math_ops.exp(
math_ops.to_float(math_ops.range(num_timescales)) * -log_timescale_increment)
scaled_time = array_ops.expand_dims(position, 1) * array_ops.expand_dims(inv_timescales, 0)
signal = array_ops.concat([math_ops.sin(scaled_time), math_ops.cos(scaled_time)], axis=1)
signal = array_ops.pad(signal, [[0, 0], [0, math_ops.mod(channels, 2)]])
if x.get_shape().ndims == 3:
signal = array_ops.reshape(signal, [1, length, channels])
else:
signal = array_ops.reshape(signal, [1, channels])
return x + signal
def _tf_range(start_or_stop, stop, step):
# Note: for static inputs (e.g. constants), tf.range errors out at graph
# construction time, instead of returning an empty tensor. Preventing the
# graph construction error aligns the semantics with Python.
# TODO(mdan): We should optimize this when a full tensor is not required.
if step is not UNDEFINED:
# TODO(mdan): Add argument coercion similar to other cases.
return math_ops.range(start_or_stop, stop, step)
if stop is not UNDEFINED:
stop = math_ops.maximum(start_or_stop, stop)
return math_ops.range(start_or_stop, stop)
start_or_stop = math_ops.maximum(start_or_stop, 0)
return math_ops.range(start_or_stop)
def dynamic_range(start_or_stop, stop=None, step=None):
"""Implementation of range using dynamic dispatch."""
if type_check.is_tensor(start_or_stop, stop, step):
if step is not None:
return math_ops.range(start_or_stop, stop, step)
if stop is not None:
return math_ops.range(start_or_stop, stop)
return math_ops.range(start_or_stop)
if step is not None:
return range(start_or_stop, stop, step)
elif stop is not None:
return range(start_or_stop, stop)
return range(start_or_stop)
def testFormatOneTensorOneDim(self):
with self.test_session():
tensor = math_ops.range(10)
format_output = string_ops.string_format("{}", tensor)
out = self.evaluate(format_output)
expected = "[0 1 2 ... 7 8 9]"
self.assertEqual(compat.as_text(out), expected)
with self.test_session():
tensor = math_ops.range(10)
format_output = string_ops.string_format("{}", [tensor])
out = self.evaluate(format_output)
expected = "[0 1 2 ... 7 8 9]"
self.assertEqual(compat.as_text(out), expected)
def _unfip_back_dims_to_front(self, x_flipped, x_shape, x_get_shape):
# E.g. suppose that originally
# chol.shape = [N1,...,Nn, k, k]
# x.shape = [M1,...,Mm, N1,...,Nn, k]
# Then we have flipped the dims so that
# x_flipped.shape = [N1,...,Nn, k, M1*...*Mm].
# We want to return x with the original shape.
rank = array_ops.rank(x_flipped)
# Permutation corresponding to [M1*...*Mm, N1,...,Nn, k]
perm = array_ops.concat(
0, (math_ops.range(rank - 1, rank), math_ops.range(0, rank - 1)))
x_with_end_at_beginning = array_ops.transpose(x_flipped, perm=perm)
x = array_ops.reshape(x_with_end_at_beginning, x_shape)
return x
def mean_pairwise_squared_error(labels,
predictions,
weights=1.0,
scope=None,
loss_collection=ops.GraphKeys.LOSSES):
"""Adds a pairwise-errors-squared loss to the training procedure.
Unlike `mean_squared_error`, which is a measure of the differences between
corresponding elements of `predictions` and `labels`,
`mean_pairwise_squared_error` is a measure of the differences between pairs of
corresponding elements of `predictions` and `labels`.
For example, if `labels`=[a, b, c] and `predictions`=[x, y, z], there are
three pairs of differences are summed to compute the loss:
loss = [ ((a-b) - (x-y)).^2 + ((a-c) - (x-z)).^2 + ((b-c) - (y-z)).^2 ] / 3
Note that since the inputs are of shape `[batch_size, d0, ... dN]`, the
corresponding pairs are computed within each batch sample but not across
samples within a batch. For example, if `predictions` represents a batch of
16 grayscale images of dimension [batch_size, 100, 200], then the set of pairs
is drawn from each image, but not across images.
`weights` acts as a coefficient for the loss. If a scalar is provided, then
the loss is simply scaled by the given value. If `weights` is a tensor of size
[batch_size], then the total loss for each sample of the batch is rescaled
by the corresponding element in the `weights` vector.
Args:
labels: The ground truth output tensor, whose shape must match the shape of
`predictions`.
predictions: The predicted outputs, a tensor of size
`[batch_size, d0, .. dN]` where N+1 is the total number of dimensions in
`predictions`.
weights: Coefficients for the loss a scalar, a tensor of shape
`[batch_size]` or a tensor whose shape matches `predictions`.
scope: The scope for the operations performed in computing the loss.
loss_collection: collection to which the loss will be added.
Returns:
A scalar `Tensor` that returns the weighted loss.
Raises:
ValueError: If the shape of `predictions` doesn't match that of `labels` or
if the shape of `weights` is invalid. Also if `labels` or `predictions
is None.
"""
if labels is None:
raise ValueError("labels must not be None.")
if predictions is None:
raise ValueError("predictions must not be None.")
with ops.name_scope(scope, "mean_pairwise_squared_error",
(predictions, labels, weights)) as scope:
weights = math_ops.to_float(weights)
labels = math_ops.to_float(labels)
with ops.control_dependencies(
(weights_broadcast_ops.assert_broadcastable(weights, labels), )):
predictions = math_ops.to_float(predictions)
predictions.get_shape().assert_is_compatible_with(
labels.get_shape())
diffs = math_ops.subtract(predictions, labels)
reduction_indices = math_ops.range(1, array_ops.rank(diffs))
sum_squares_diff_per_batch = math_ops.reduce_sum(
math_ops.square(diffs),
reduction_indices=reduction_indices,
keep_dims=True)
num_present_per_batch = _num_present(diffs,
weights,
per_batch=True)
term1 = 2.0 * _safe_div(sum_squares_diff_per_batch,
num_present_per_batch)
sum_diff = math_ops.reduce_sum(diffs,
reduction_indices=reduction_indices,
keep_dims=True)
term2 = 2.0 * _safe_div(math_ops.square(sum_diff),
math_ops.square(num_present_per_batch))
weighted_losses = math_ops.multiply(term1 - term2, weights)
loss = math_ops.reduce_sum(weighted_losses)
mean_loss = array_ops.where(
math_ops.reduce_sum(num_present_per_batch) > 0,
loss,
array_ops.zeros_like(loss),
name="value")
util.add_loss(mean_loss, loss_collection)
return mean_loss
class RaggedMapOpTest(ragged_test_util.RaggedTensorTestCase,
parameterized.TestCase):
@parameterized.parameters([
# The following test sets map over a RaggedTensor and apply a
# transformation that returns with shape:
# [d1, (d2)] -> [d1]
dict(
fn=mo.reduce_mean,
elems=[[1, 2, 3], [4, 5], [6, 7]],
expected_output=[2, 4, 6],
),
dict(
fn=string_ops.reduce_join,
elems=[['foo', 'bar', 'baz'], ['a'], ['b', 'c']],
expected_output=[b'foobarbaz', b'a', b'bc'],
dtype=dtypes.string,
),
# [d1, (d2)] -> [d1, 2]
dict(
fn=lambda x: array_ops.stack([mo.reduce_mean(x),
mo.reduce_sum(x)]),
# fn=self.stack_mean_and_sum,
elems=[[1, 2, 3], [4, 5], [6, 7]],
expected_output=[[2, 6], [4.5, 9], [6.5, 13]],
dtype=dtypes.float32,
expected_ragged_rank=0,
),
# [d1, (d2)] -> [d1, (d2)]
dict(
fn=lambda x: x + np.int64(1),
elems=[[1, 2, 3], [4, 5], [6, 7]],
expected_output=[[2, 3, 4], [5, 6], [7, 8]],
dtype=dtypes.int64,
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=1),
),
# [d1, (d2), d3] -> [d1, (d2), d3]
dict(
fn=lambda x: x + np.int64(1),
elems=[[[1, 2], [3, 4]], [], [[5, 6], [7, 8], [9, 0]]],
elems_ragged_rank=1,
expected_ragged_rank=1,
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=1),
expected_output=[[[2, 3], [4, 5]], [], [[6, 7], [8, 9], [10, 1]]],
),
# [d1, (d2)] -> [d1, (d2), (d3)]
dict(
fn=lambda x: ragged_tensor.RaggedTensor.from_row_starts(x, [0]),
elems=[[1, 2, 3], [4, 5], [6, 7]],
expected_output=[[[1, 2, 3]], [[4, 5]], [[6, 7]]],
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=2),
),
# [d1, (d2), (d3)] -> [d1, (d2), (d3)]
dict(
fn=lambda x: ragged_functional_ops.map_flat_values(mo.add, x, 1),
elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
expected_output=[[[2, 3, 4]], [[5, 6], [7, 8]]],
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=2),
),
# [d1, (d2), (d3)] -> [d1, (d2)]
dict(
fn=lambda x: ragged_math_ops.reduce_sum(x, axis=1),
elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
expected_output=[[6], [9, 13]],
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=1),
),
# [d1, (d2), (d3)] -> [d1, (d3)]
dict(
fn=lambda x: ragged_math_ops.reduce_sum(x, axis=0),
elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
expected_output=[[1, 2, 3], [10, 12]],
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=1),
),
# [d1, (d2), (d3)] -> [d1]
dict(
fn=ragged_math_ops.reduce_sum,
elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
expected_output=[6, 22],
result_dtype=dtypes.int64,
),
# [d1] -> [d1, (d2)]
dict(
fn=mo.range,
elems=[4, 0, 2],
expected_output=[[0, 1, 2, 3], [], [0, 1]],
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=1),
),
# [d1] -> [d1, (d2), (d3)]
dict(
fn=lambda x: ragged_math_ops.range(mo.range(x)),
elems=[5, 0, 3],
expected_output=[[[], [0], [0, 1], [0, 1, 2], [0, 1, 2, 3]], [],
[[], [0], [0, 1]]],
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=2),
),
# [d1, (d2), (d3), (d4a), (d5)] -> [d1, (d2), (d3), (d4b), (d5)]
dict(
fn=lambda x: x + np.int64(1),
elems=[[[[[1, 2, 3]], [[4], [5]]]], [[[[6, 7]]], [[[8], []]]]],
expected_output=[[[[[2, 3, 4]], [[5], [6]]]],
[[[[7, 8]]], [[[9], []]]]],
result_dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64,
ragged_rank=4),
),
])
def testRaggedMap(
self,
fn,
elems,
expected_output,
expected_ragged_rank=None,
result_ragged_rank=None,
elems_ragged_rank=None,
dtype=dtypes.int64,
result_dtype=None,
infer_shape=False,
):
elems = ragged_factory_ops.constant(elems, dtype, elems_ragged_rank)
output = ragged_map_ops.map_fn(fn=fn,
elems=elems,
dtype=result_dtype,
infer_shape=infer_shape)
expected_rt = ragged_factory_ops.constant(
expected_output, ragged_rank=expected_ragged_rank)
self.assertRaggedEqual(expected_rt, output)
def testRaggedMapOnStructure(self):
batman = ragged_factory_ops.constant([[1, 2, 3], [4], [5, 6, 7]])
# [[10, 20, 30], [40], [50, 60, 70]]
robin = ragged_functional_ops.map_flat_values(mo.multiply, batman, 10)
features = {'batman': batman, 'robin': robin}
def _reduce_sum_from_all(f):
return mo.reduce_sum(f['batman']) + mo.reduce_sum(f['robin'])
output = ragged_map_ops.map_fn(
fn=_reduce_sum_from_all,
elems=features,
dtype=dtypes.int32,
)
self.assertRaggedEqual(output, [66, 44, 198])
# Test mapping over a dict of RTs can produce a dict of RTs.
def testRaggedMapOnStructure_RaggedOutputs(self):
batman = ragged_factory_ops.constant([[1, 2, 3], [4], [5, 6, 7]])
# [[10, 20, 30], [40], [50, 60, 70]]
robin = ragged_functional_ops.map_flat_values(mo.multiply, batman, 10)
features = {'batman': batman, 'robin': robin}
def _increment(f):
return {
'batman': f['batman'] + 1,
'robin': f['robin'] + 1,
}
output = ragged_map_ops.map_fn(
fn=_increment,
elems=features,
infer_shape=False,
dtype={
'batman':
ragged_tensor.RaggedTensorType(dtype=dtypes.int32,
ragged_rank=1),
'robin':
ragged_tensor.RaggedTensorType(dtype=dtypes.int32,
ragged_rank=1)
},
)
self.assertRaggedEqual(output['batman'], [[2, 3, 4], [5], [6, 7, 8]])
self.assertRaggedEqual(output['robin'],
[[11, 21, 31], [41], [51, 61, 71]])
def testZip(self):
x = ragged_factory_ops.constant(
[[10, 20], [30, 40], [50, 60], [70], [80, 90, 100]], dtypes.int64)
y = array_ops.expand_dims(mo.range(x.nrows(out_type=dtypes.int64)),
axis=1)
def _zip(foo):
y_val, x_val = foo
bar = backend.tile(y_val, array_ops.shape(x_val))
return array_ops.stack([bar, x_val], axis=1)
output = ragged_map_ops.map_fn(_zip, (y, x),
dtype=ragged_tensor.RaggedTensorType(
dtype=dtypes.int64, ragged_rank=1),
infer_shape=False)
self.assertRaggedEqual(
output,
[[[0, 10], [0, 20]], [[1, 30], [1, 40]], [[2, 50], [2, 60]],
[[3, 70]], [[4, 80], [4, 90], [4, 100]]])
def testBatchGather(self):
tokens = ragged_factory_ops.constant([['hello', '.', 'there'],
['merhaba'],
['bonjour', '.', 'ca va', '?']])
indices = ragged_factory_ops.constant([[0, 2], [0], [0, 2]])
def gather(x):
tokens_val, indices_val = x
return array_ops.gather(tokens_val, indices_val)
data = tokens, indices
out = ragged_map_ops.map_fn(gather,
data,
dtype=ragged_tensor.RaggedTensorType(
dtype=dtypes.string, ragged_rank=1),
infer_shape=False)
self.assertRaggedEqual(
out, [[b'hello', b'there'], [b'merhaba'], [b'bonjour', b'ca va']])
def testMismatchRaggedRank(self):
elems = ragged_factory_ops.constant([[[1, 2, 3]], [[4, 5], [6, 7]]])
fn = lambda x: ragged_math_ops.reduce_sum(x, axis=0)
with self.assertRaisesWithLiteralMatch(
ValueError,
r'The declared ragged rank (23) mismatches the result (1)'):
_ = ragged_map_ops.map_fn(fn,
elems,
dtype=ragged_tensor.RaggedTensorType(
dtype=dtypes.int64, ragged_rank=23))
def testMismatchRaggedRank2(self):
elems = ragged_factory_ops.constant([[1, 2, 3], [4, 5], [6, 7]])
fn = lambda x: ragged_tensor.RaggedTensor.from_row_starts(x, [0])
with self.assertRaisesWithLiteralMatch(
ValueError,
r'The declared ragged rank (10) mismatches the result (2)'):
_ = ragged_map_ops.map_fn(fn,
elems,
dtype=ragged_tensor.RaggedTensorType(
dtype=dtypes.int64, ragged_rank=10))
def testMapOnSparseTensor(self):
s = sparse_tensor.SparseTensor(
indices=[[0, 0], [0, 1], [1, 0], [1, 1]],
values=[0, 5, 0, 4],
dense_shape=[2, 2],
)
t2 = ragged_tensor.RaggedTensor.from_sparse(s)
id_t2 = ragged_map_ops.map_fn(
lambda x: x,
t2,
)
self.assertRaggedEqual(id_t2, [[0, 5], [0, 4]])
def _to_dense(self):
perm = ops.convert_to_tensor_v2_with_dispatch(self.perm)
return math_ops.cast(
math_ops.equal(
math_ops.range(0, self._domain_dimension_tensor(perm)),
perm[..., array_ops.newaxis]), self.dtype)
def _diag_part(self):
perm = ops.convert_to_tensor_v2_with_dispatch(self.perm)
return math_ops.cast(
math_ops.equal(
math_ops.range(0, self._domain_dimension_tensor(perm)), perm),
self.dtype)
def unstack(self, value, name=None):
"""See TensorArray."""
with ops.name_scope(name, "TensorArrayUnstack", [self._handle, value]):
num_elements = array_ops.shape(value)[0]
return self.scatter(
indices=math_ops.range(0, num_elements), value=value, name=name)
def norm(tensor,
ord='euclidean',
axis=None,
keepdims=None,
name=None,
keep_dims=None):
r"""Computes the norm of vectors, matrices, and tensors.
This function can compute several different vector norms (the 1-norm, the
Euclidean or 2-norm, the inf-norm, and in general the p-norm for p > 0) and
matrix norms (Frobenius, 1-norm, 2-norm and inf-norm).
Args:
tensor: `Tensor` of types `float32`, `float64`, `complex64`, `complex128`
ord: Order of the norm. Supported values are 'fro', 'euclidean',
`1`, `2`, `np.inf` and any positive real number yielding the corresponding
p-norm. Default is 'euclidean' which is equivalent to Frobenius norm if
`tensor` is a matrix and equivalent to 2-norm for vectors.
Some restrictions apply:
a) The Frobenius norm `fro` is not defined for vectors,
b) If axis is a 2-tuple (matrix norm), only 'euclidean', 'fro', `1`,
`2`, `np.inf` are supported.
See the description of `axis` on how to compute norms for a batch of
vectors or matrices stored in a tensor.
axis: If `axis` is `None` (the default), the input is considered a vector
and a single vector norm is computed over the entire set of values in the
tensor, i.e. `norm(tensor, ord=ord)` is equivalent to
`norm(reshape(tensor, [-1]), ord=ord)`.
If `axis` is a Python integer, the input is considered a batch of vectors,
and `axis` determines the axis in `tensor` over which to compute vector
norms.
If `axis` is a 2-tuple of Python integers it is considered a batch of
matrices and `axis` determines the axes in `tensor` over which to compute
a matrix norm.
Negative indices are supported. Example: If you are passing a tensor that
can be either a matrix or a batch of matrices at runtime, pass
`axis=[-2,-1]` instead of `axis=None` to make sure that matrix norms are
computed.
keepdims: If True, the axis indicated in `axis` are kept with size 1.
Otherwise, the dimensions in `axis` are removed from the output shape.
name: The name of the op.
keep_dims: Deprecated alias for `keepdims`.
Returns:
output: A `Tensor` of the same type as tensor, containing the vector or
matrix norms. If `keepdims` is True then the rank of output is equal to
the rank of `tensor`. Otherwise, if `axis` is none the output is a scalar,
if `axis` is an integer, the rank of `output` is one less than the rank
of `tensor`, if `axis` is a 2-tuple the rank of `output` is two less
than the rank of `tensor`.
Raises:
ValueError: If `ord` or `axis` is invalid.
@compatibility(numpy)
Mostly equivalent to numpy.linalg.norm.
Not supported: ord <= 0, 2-norm for matrices, nuclear norm.
Other differences:
a) If axis is `None`, treats the flattened `tensor` as a vector
regardless of rank.
b) Explicitly supports 'euclidean' norm as the default, including for
higher order tensors.
@end_compatibility
"""
keepdims = deprecation.deprecated_argument_lookup('keepdims', keepdims,
'keep_dims', keep_dims)
if keepdims is None:
keepdims = False
is_matrix_norm = ((isinstance(axis, tuple) or isinstance(axis, list)) and
len(axis) == 2)
if is_matrix_norm:
axis = tuple(axis)
if (not isinstance(axis[0], int) or not isinstance(axis[1], int) or
axis[0] == axis[1]):
raise ValueError(
"'axis' must be None, an integer, or a tuple of 2 "
f"unique integers, got {axis}")
supported_matrix_norms = ['euclidean', 'fro', 1, 2, np.inf]
if ord not in supported_matrix_norms:
raise ValueError(f"'ord' must be a supported matrix norm in "
f"{supported_matrix_norms}, got {ord}")
else:
if not (isinstance(axis, int) or axis is None):
raise ValueError(
"'axis' must be None, an integer, or a "
f"tuple of 2 unique integers, got {axis}")
supported_vector_norms = ['euclidean', 1, 2, np.inf]
if (not np.isreal(ord) or ord <= 0) and ord not in supported_vector_norms:
raise ValueError(f"'ord' must be a supported vector norm, got {ord}")
if axis is not None:
axis = (axis,)
with ops.name_scope(name, 'norm', [tensor]):
tensor = ops.convert_to_tensor(tensor)
if ord in ['fro', 'euclidean', 2, 2.0]:
if is_matrix_norm and ord in [2, 2.0]:
rank = array_ops.rank(tensor)
positive_axis = map_fn.map_fn(
lambda i: control_flow_ops.cond(i >= 0, lambda: i, lambda: i + rank
), ops.convert_to_tensor(axis))
axes = math_ops.range(rank)
perm_before = array_ops.concat([
gen_array_ops.list_diff(axes, positive_axis, dtypes.int32)[0],
positive_axis
],
axis=0)
perm_after = map_fn.map_fn(
lambda i: math_ops.cast(
array_ops.squeeze(
array_ops.where_v2(math_ops.equal(perm_before, i))),
dtype=dtypes.int32), axes)
permed = array_ops.transpose(tensor, perm=perm_before)
matrix_2_norm = array_ops.expand_dims(
math_ops.reduce_max(
math_ops.abs(gen_linalg_ops.svd(permed, compute_uv=False)[0]),
axis=-1,
keepdims=True),
axis=-1)
result = array_ops.transpose(matrix_2_norm, perm=perm_after)
else:
result = math_ops.sqrt(
math_ops.reduce_sum(
tensor * math_ops.conj(tensor), axis, keepdims=True))
# TODO(rmlarsen): Replace with the following, once gradients are defined
# result = math_ops.reduce_euclidean_norm(tensor, axis, keepdims=True)
else:
result = math_ops.abs(tensor)
if ord == 1:
sum_axis = None if axis is None else axis[0]
result = math_ops.reduce_sum(result, sum_axis, keepdims=True)
if is_matrix_norm:
result = math_ops.reduce_max(result, axis[-1], keepdims=True)
elif ord == np.inf:
if is_matrix_norm:
result = math_ops.reduce_sum(result, axis[1], keepdims=True)
max_axis = None if axis is None else axis[0]
result = math_ops.reduce_max(result, max_axis, keepdims=True)
else:
# General p-norms (positive p only)
result = math_ops.pow(
math_ops.reduce_sum(math_ops.pow(result, ord), axis, keepdims=True),
1.0 / ord)
if not keepdims:
result = array_ops.squeeze(result, axis)
return result
def sparse_fill_empty_rows(sp_input, default_value, name=None):
"""Fills empty rows in the input 2-D `SparseTensor` with a default value.
This op adds entries with the specified `default_value` at index
`[row, 0]` for any row in the input that does not already have a value.
For example, suppose `sp_input` has shape `[5, 6]` and non-empty values:
[0, 1]: a
[0, 3]: b
[2, 0]: c
[3, 1]: d
Rows 1 and 4 are empty, so the output will be of shape `[5, 6]` with values:
[0, 1]: a
[0, 3]: b
[1, 0]: default_value
[2, 0]: c
[3, 1]: d
[4, 0]: default_value
Note that the input may have empty columns at the end, with no effect on
this op.
The output `SparseTensor` will be in row-major order and will have the
same shape as the input.
This op also returns an indicator vector such that
empty_row_indicator[i] = True iff row i was an empty row.
Args:
sp_input: A `SparseTensor` with shape `[N, M]`.
default_value: The value to fill for empty rows, with the same type as
`sp_input.`
name: A name prefix for the returned tensors (optional)
Returns:
sp_ordered_output: A `SparseTensor` with shape `[N, M]`, and with all empty
rows filled in with `default_value`.
empty_row_indicator: A bool vector of length `N` indicating whether each
input row was empty.
Raises:
TypeError: If `sp_input` is not a `SparseTensor`.
"""
if not isinstance(sp_input, ops.SparseTensor):
raise TypeError("Input must be a SparseTensor")
with ops.op_scope([sp_input], name, "SparseFillEmptyRows"):
default_value = ops.convert_to_tensor(default_value,
dtype=sp_input.values.dtype)
num_rows = math_ops.cast(sp_input.shape[0], dtypes.int32)
all_row_indices = math_ops.cast(math_ops.range(num_rows), dtypes.int64)
empty_row_indices, _ = array_ops.list_diff(all_row_indices,
sp_input.indices[:, 0])
empty_row_indicator = sparse_to_dense(
empty_row_indices, array_ops.expand_dims(sp_input.shape[0], -1), True,
False)
empty_row_indices_as_column = array_ops.reshape(empty_row_indices, [-1, 1])
additional_indices = array_ops.concat(
1, [empty_row_indices_as_column,
array_ops.zeros_like(empty_row_indices_as_column)])
additional_values = array_ops.fill(
array_ops.shape(empty_row_indices), default_value)
all_indices_unordered = array_ops.concat(0, [sp_input.indices,
additional_indices])
all_values_unordered = array_ops.concat(0, [sp_input.values,
additional_values])
sp_unordered_output = ops.SparseTensor(all_indices_unordered,
all_values_unordered, sp_input.shape)
sp_ordered_output = sparse_reorder(sp_unordered_output)
return sp_ordered_output, empty_row_indicator
def _sample_n(self, n, seed=None):
with ops.control_dependencies(self._assertions):
n = ops.convert_to_tensor(n, name="n")
static_n = tensor_util.constant_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.get_shape()
if static_samples_shape.is_fully_defined():
samples_shape = static_samples_shape.as_list()
samples_size = static_samples_shape.num_elements()
else:
samples_shape = array_ops.shape(cat_samples)
samples_size = array_ops.size(cat_samples)
static_batch_shape = self.batch_shape
if static_batch_shape.is_fully_defined():
batch_shape = static_batch_shape.as_list()
batch_size = static_batch_shape.num_elements()
else:
batch_shape = self.batch_shape_tensor()
batch_size = math_ops.reduce_prod(batch_shape)
static_event_shape = self.event_shape
if static_event_shape.is_fully_defined():
event_shape = np.array(static_event_shape.as_list(),
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 = array_ops.reshape(
math_ops.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 = data_flow_ops.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 = array_ops.reshape(
array_ops.tile(math_ops.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 = data_flow_ops.dynamic_partition(
data=batch_raw_indices,
partitions=cat_samples,
num_partitions=self.num_components)
samples_class = [None for _ in range(self.num_components)]
for c in range(self.num_components):
n_class = array_ops.size(partitioned_samples_indices[c])
seed = distribution_util.gen_new_seed(seed, "mixture")
samples_class_c = self.components[c].sample(n_class, seed=seed)
# 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 * math_ops.range(n_class) +
partitioned_batch_indices[c])
samples_class_c = array_ops.reshape(
samples_class_c,
array_ops.concat([[n_class * batch_size], event_shape], 0))
samples_class_c = array_ops.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 = data_flow_ops.dynamic_stitch(
indices=partitioned_samples_indices, data=samples_class)
# Reshape back to proper sample, batch, and event shape.
ret = array_ops.reshape(
lhs_flat_ret,
array_ops.concat(
[samples_shape, self.event_shape_tensor()], 0))
ret.set_shape(
tensor_shape.TensorShape(static_samples_shape).concatenate(
self.event_shape))
return ret
def testDType(self):
zero_int32 = math_ops.cast(0, dtypes.int32)
zero_int64 = math_ops.cast(0, dtypes.int64)
zero_float32 = math_ops.cast(0, dtypes.float32)
zero_float64 = math_ops.cast(0, dtypes.float64)
self.assertEqual(math_ops.range(zero_int32, 0, 1).dtype, dtypes.int32)
self.assertEqual(math_ops.range(zero_int64, 0, 1).dtype, dtypes.int64)
self.assertEqual(math_ops.range(zero_float32, 0, 1).dtype, dtypes.float32)
self.assertEqual(math_ops.range(zero_float64, 0, 1).dtype, dtypes.float64)
self.assertEqual(
math_ops.range(zero_int32, zero_int64, 1).dtype, dtypes.int64)
self.assertEqual(
math_ops.range(zero_int64, zero_float32, 1).dtype, dtypes.float32)
self.assertEqual(
math_ops.range(zero_float32, zero_float64, 1).dtype, dtypes.float64)
self.assertEqual(
math_ops.range(zero_float64, zero_int32, 1).dtype, dtypes.float64)
self.assertEqual(
math_ops.range(
0, 0, 1, dtype=dtypes.int32).dtype, dtypes.int32)
self.assertEqual(
math_ops.range(
0, 0, 1, dtype=dtypes.int64).dtype, dtypes.int64)
self.assertEqual(
math_ops.range(
0, 0, 1, dtype=dtypes.float32).dtype, dtypes.float32)
self.assertEqual(
math_ops.range(
0, 0, 1, dtype=dtypes.float64).dtype, dtypes.float64)
def train_fn(self, iterator):
for _ in math_ops.range(5):
x = math_ops.matmul(array_ops.squeeze(next(iterator)), self.w)
x = math_ops.matmul(random_ops.random_uniform((1000, 1000)), x)
self.w.assign_add(x)
self.iterations.assign_add(1)
def _sparse(i):
return sparse_tensor.SparseTensorValue(
indices=array_ops.expand_dims(
math_ops.range(i, dtype=dtypes.int64), 1),
values=array_ops.fill([math_ops.cast(i, dtypes.int32)], i),
dense_shape=[i])
def rotate_transpose(x, shift, name="rotate_transpose"):
"""Circularly moves dims left or right.
Effectively identical to:
```python
numpy.transpose(x, numpy.roll(numpy.arange(len(x.shape)), shift))
```
When `validate_args=True` additional graph-runtime checks are
performed. These checks entail moving data from to GPU to CPU.
Example:
```python
x = ... # Tensor of shape [1, 2, 3, 4].
rotate_transpose(x, -1) # result shape: [2, 3, 4, 1]
rotate_transpose(x, -2) # result shape: [3, 4, 1, 2]
rotate_transpose(x, 1) # result shape: [4, 1, 2, 3]
rotate_transpose(x, 2) # result shape: [3, 4, 1, 2]
rotate_transpose(x, 7) == rotate_transpose(x, 3)
rotate_transpose(x, -7) == rotate_transpose(x, -3)
```
Args:
x: `Tensor`.
shift: `Tensor`. Number of dimensions to transpose left (shift<0) or
transpose right (shift>0).
name: `String`. The name to give this op.
Returns:
rotated_x: Input `Tensor` with dimensions circularly rotated by shift.
Raises:
TypeError: if shift is not integer type.
"""
with ops.name_scope(name, values=[x, shift]):
x = ops.convert_to_tensor(x, name="x")
shift = ops.convert_to_tensor(shift, name="shift")
# We do not assign back to preserve constant-ness.
check_ops.assert_integer(shift)
shift_value_static = tensor_util.constant_value(shift)
ndims = x.get_shape().ndims
if ndims is not None and shift_value_static is not None:
if ndims < 2: return x
shift_value_static = np.sign(shift_value_static) * (
abs(shift_value_static) % ndims)
if shift_value_static == 0: return x
perm = np.roll(np.arange(ndims), shift_value_static)
return array_ops.transpose(x, perm=perm)
else:
# Consider if we always had a positive shift, and some specified
# direction.
# When shifting left we want the new array:
# last(x, n-shift) + first(x, shift)
# and if shifting right then we want:
# last(x, shift) + first(x, n-shift)
# Observe that last(a) == slice(a, n) and first(a) == slice(0, a).
# Also, we can encode direction and shift as one: direction * shift.
# Combining these facts, we have:
# a = cond(shift<0, -shift, n-shift)
# last(x, n-a) + first(x, a) == x[a:n] + x[0:a]
# Finally, we transform shift by modulo length so it can be specified
# independently from the array upon which it operates (like python).
ndims = array_ops.rank(x)
shift = math_ops.select(math_ops.less(shift, 0),
math_ops.mod(-shift, ndims),
ndims - math_ops.mod(shift, ndims))
first = math_ops.range(0, shift)
last = math_ops.range(shift, ndims)
perm = array_ops.concat(0, (last, first))
return array_ops.transpose(x, perm=perm)
def sample_chain(num_results,
target_log_prob_fn,
current_state,
step_size,
num_leapfrog_steps,
num_burnin_steps=0,
num_steps_between_results=0,
seed=None,
current_target_log_prob=None,
current_grads_target_log_prob=None,
name=None):
"""Runs multiple iterations of one or more Hamiltonian Monte Carlo chains.
Hamiltonian Monte Carlo (HMC) is a Markov chain Monte Carlo (MCMC) algorithm
that takes a series of gradient-informed steps to produce a Metropolis
proposal. This function samples from an HMC Markov chain at `current_state`
and whose stationary distribution has log-unnormalized-density
`target_log_prob_fn()`.
This function samples from multiple chains in parallel. It assumes that the
the leftmost dimensions of (each) `current_state` (part) index an independent
chain. The function `target_log_prob_fn()` sums log-probabilities across
event dimensions (i.e., current state (part) rightmost dimensions). Each
element of the output of `target_log_prob_fn()` represents the (possibly
unnormalized) log-probability of the joint distribution over (all) the current
state (parts).
The `current_state` can be represented as a single `Tensor` or a `list` of
`Tensors` which collectively represent the current state. When specifying a
`list`, one must also specify a list of `step_size`s.
Note: `target_log_prob_fn` is called exactly twice.
Only one out of every `num_steps_between_samples + 1` steps is included in the
returned results. This "thinning" comes at a cost of reduced statistical
power, while reducing memory requirements and autocorrelation. For more
discussion see [1].
[1]: "Statistically efficient thinning of a Markov chain sampler."
Art B. Owen. April 2017.
http://statweb.stanford.edu/~owen/reports/bestthinning.pdf
#### Examples:
##### Sample from a diagonal-variance Gaussian.
```python
tfd = tf.contrib.distributions
def make_likelihood(true_variances):
return tfd.MultivariateNormalDiag(
scale_diag=tf.sqrt(true_variances))
dims = 10
dtype = np.float32
true_variances = tf.linspace(dtype(1), dtype(3), dims)
likelihood = make_likelihood(true_variances)
states, kernel_results = hmc.sample_chain(
num_results=1000,
target_log_prob_fn=likelihood.log_prob,
current_state=tf.zeros(dims),
step_size=0.5,
num_leapfrog_steps=2,
num_burnin_steps=500)
# Compute sample stats.
sample_mean = tf.reduce_mean(states, axis=0)
sample_var = tf.reduce_mean(
tf.squared_difference(states, sample_mean),
axis=0)
```
##### Sampling from factor-analysis posteriors with known factors.
I.e.,
```none
for i=1..n:
w[i] ~ Normal(0, eye(d)) # prior
x[i] ~ Normal(loc=matmul(w[i], F)) # likelihood
```
where `F` denotes factors.
```python
tfd = tf.contrib.distributions
def make_prior(dims, dtype):
return tfd.MultivariateNormalDiag(
loc=tf.zeros(dims, dtype))
def make_likelihood(weights, factors):
return tfd.MultivariateNormalDiag(
loc=tf.tensordot(weights, factors, axes=[[0], [-1]]))
# Setup data.
num_weights = 10
num_factors = 4
num_chains = 100
dtype = np.float32
prior = make_prior(num_weights, dtype)
weights = prior.sample(num_chains)
factors = np.random.randn(num_factors, num_weights).astype(dtype)
x = make_likelihood(weights, factors).sample(num_chains)
def target_log_prob(w):
# Target joint is: `f(w) = p(w, x | factors)`.
return prior.log_prob(w) + make_likelihood(w, factors).log_prob(x)
# Get `num_results` samples from `num_chains` independent chains.
chains_states, kernels_results = hmc.sample_chain(
num_results=1000,
target_log_prob_fn=target_log_prob,
current_state=tf.zeros([num_chains, dims], dtype),
step_size=0.1,
num_leapfrog_steps=2,
num_burnin_steps=500)
# Compute sample stats.
sample_mean = tf.reduce_mean(chains_states, axis=[0, 1])
sample_var = tf.reduce_mean(
tf.squared_difference(chains_states, sample_mean),
axis=[0, 1])
```
Args:
num_results: Integer number of Markov chain draws.
target_log_prob_fn: Python callable which takes an argument like
`current_state` (or `*current_state` if it's a list) and returns its
(possibly unnormalized) log-density under the target distribution.
current_state: `Tensor` or Python `list` of `Tensor`s representing the
current state(s) of the Markov chain(s). The first `r` dimensions index
independent chains, `r = tf.rank(target_log_prob_fn(*current_state))`.
step_size: `Tensor` or Python `list` of `Tensor`s representing the step size
for the leapfrog integrator. Must broadcast with the shape of
`current_state`. Larger step sizes lead to faster progress, but too-large
step sizes make rejection exponentially more likely. When possible, it's
often helpful to match per-variable step sizes to the standard deviations
of the target distribution in each variable.
num_leapfrog_steps: Integer number of steps to run the leapfrog integrator
for. Total progress per HMC step is roughly proportional to `step_size *
num_leapfrog_steps`.
num_burnin_steps: Integer number of chain steps to take before starting to
collect results.
Default value: 0 (i.e., no burn-in).
num_steps_between_results: Integer number of chain steps between collecting
a result. Only one out of every `num_steps_between_samples + 1` steps is
included in the returned results. This "thinning" comes at a cost of
reduced statistical power, while reducing memory requirements and
autocorrelation. For more discussion see [1].
Default value: 0 (i.e., no subsampling).
seed: Python integer to seed the random number generator.
current_target_log_prob: (Optional) `Tensor` representing the value of
`target_log_prob_fn` at the `current_state`. The only reason to specify
this argument is to reduce TF graph size.
Default value: `None` (i.e., compute as needed).
current_grads_target_log_prob: (Optional) Python list of `Tensor`s
representing gradient of `target_log_prob` at the `current_state` and wrt
the `current_state`. Must have same shape as `current_state`. The only
reason to specify this argument is to reduce TF graph size.
Default value: `None` (i.e., compute as needed).
name: Python `str` name prefixed to Ops created by this function.
Default value: `None` (i.e., "hmc_sample_chain").
Returns:
accepted_states: Tensor or Python list of `Tensor`s representing the
state(s) of the Markov chain(s) at each result step. Has same shape as
input `current_state` but with a prepended `num_results`-size dimension.
kernel_results: `collections.namedtuple` of internal calculations used to
advance the chain.
"""
with ops.name_scope(name, "hmc_sample_chain", [
num_results, current_state, step_size, num_leapfrog_steps,
num_burnin_steps, num_steps_between_results, seed,
current_target_log_prob, current_grads_target_log_prob
]):
with ops.name_scope("initialize"):
[
current_state,
step_size,
current_target_log_prob,
current_grads_target_log_prob,
] = _prepare_args(target_log_prob_fn, current_state, step_size,
current_target_log_prob,
current_grads_target_log_prob)
num_results = ops.convert_to_tensor(num_results,
dtype=dtypes.int32,
name="num_results")
num_leapfrog_steps = ops.convert_to_tensor(
num_leapfrog_steps,
dtype=dtypes.int32,
name="num_leapfrog_steps")
num_burnin_steps = ops.convert_to_tensor(num_burnin_steps,
dtype=dtypes.int32,
name="num_burnin_steps")
num_steps_between_results = ops.convert_to_tensor(
num_steps_between_results,
dtype=dtypes.int32,
name="num_steps_between_results")
def _run_chain(num_steps, current_state, kernel_results):
"""Runs the chain(s) for `num_steps`."""
def _loop_body(iter_, current_state, kernel_results):
return [iter_ + 1] + list(
kernel(target_log_prob_fn, current_state, step_size,
num_leapfrog_steps, seed,
kernel_results.current_target_log_prob,
kernel_results.current_grads_target_log_prob))
while_loop_kwargs = dict(
cond=lambda iter_, *args: iter_ < num_steps,
body=_loop_body,
loop_vars=[
np.int32(0),
current_state,
kernel_results,
],
)
if seed is not None:
while_loop_kwargs["parallel_iterations"] = 1
return control_flow_ops.while_loop(
**while_loop_kwargs)[1:] # Lop-off "iter_".
def _scan_body(args_list, iter_):
"""Closure which implements `tf.scan` body."""
current_state, kernel_results = args_list
return _run_chain(
1 + array_ops.where(math_ops.equal(iter_, 0), num_burnin_steps,
num_steps_between_results), current_state,
kernel_results)
scan_kwargs = dict(
fn=_scan_body,
elems=math_ops.range(num_results), # iter_: used to choose burnin.
initializer=[
current_state,
_make_dummy_kernel_results(current_state,
current_target_log_prob,
current_grads_target_log_prob),
])
if seed is not None:
scan_kwargs["parallel_iterations"] = 1
return functional_ops.scan(**scan_kwargs)
def _embedding_lookup_and_transform(params,
ids,
partition_strategy="mod",
name=None,
max_norm=None,
transform_fn=None):
"""Helper function for embedding_lookup and _compute_sampled_logits.
This function is a generalization of embedding_lookup that optionally
applies a caller-specified transformation to each embedding. This is
done through the `transform_fn` argument. If provided, the function is
applied to each partitioned tensor of retrieved embeddings, colocated
with the embeddings. This function will be called with a single `Tensor`
argument of the same type as the `params` tensor and should return a
`Tensor`. The shape of the argument will be the same as `params` except
for the size of the first dimension. The first dimension of the result's
shape must be the same size as the argument's.
Args:
params: See embedding_lookup.
ids: See embedding_lookup.
partition_strategy: See embedding_lookup.
name: See embedding_lookup.
max_norm: See embedding_lookup.
transform_fn: An optional function to apply to each retrieved embedding. If
max_norm is provided, transform_fn is applied to the norm-limited
embeddings.
Returns:
See embedding_lookup for details.
Raises:
ValueError: If `params` is empty.
"""
if params is None:
raise ValueError("params must be specified")
if isinstance(params, (list, tuple)) and not params:
raise ValueError("Length of params is currently 0. "
"Need at least one param.")
if isinstance(params, variables.PartitionedVariable):
params = list(params) # Iterate to get the underlying Variables.
if not isinstance(params, list):
params = [params]
with ops.name_scope(name, "embedding_lookup", params + [ids]) as name:
np = len(params) # Number of partitions
# Preserve the resource variable status to avoid accidental dense reads.
if not any(
isinstance(p, resource_variable_ops.BaseResourceVariable)
for p in params):
params = ops.convert_n_to_tensor_or_indexed_slices(params,
name="params")
ids = ops.convert_to_tensor(ids, name="ids")
if np == 1 and (not transform_fn or ids.get_shape().ndims == 1):
with _colocate_with(params[0]):
result = _clip(array_ops.gather(params[0], ids, name=name),
ids, max_norm)
if transform_fn:
result = transform_fn(result)
# Make sure the final result does not have colocation constraints on the
# params. Similar to the case np > 1 where parallel_dynamic_stitch is
# outside the scope of all with _colocate_with(params[p]).
return array_ops.identity(result)
else:
# Flatten the ids. There are two cases where we need to do this.
# - There is more than one params tensor.
# - There is a transform_fn and ids is not statically known to be 1-D.
# We must flatten in this case because transform_fn expects a flat
# tensor of embeddings.
flat_ids = array_ops.reshape(ids, [-1])
original_indices = math_ops.range(array_ops.size(flat_ids))
# Create p_assignments and set new_ids depending on the strategy.
if partition_strategy == "mod":
p_assignments = flat_ids % np
new_ids = flat_ids // np
elif partition_strategy == "div":
# Compute num_total_ids as the sum of dim-0 of params, then assign to
# partitions based on a constant number of ids per partition. Optimize
# if we already know the full shape statically.
dim_0_size = tensor_shape.Dimension(
tensor_shape.dimension_value(params[0].get_shape()[0]))
for p in range(1, np):
dim_0_size += tensor_shape.Dimension(
tensor_shape.dimension_value(params[p].get_shape()[0]))
if dim_0_size.value:
num_total_ids = constant_op.constant(
dim_0_size.value, flat_ids.dtype)
else:
dim_0_sizes = []
for p in range(np):
param_p_dim = tensor_shape.dimension_value(
params[p].get_shape()[0])
if param_p_dim is not None:
dim_0_sizes.append(param_p_dim)
else:
with _colocate_with(params[p]):
dim_0_sizes.append(
array_ops.shape(params[p])[0])
num_total_ids = math_ops.reduce_sum(
math_ops.cast(array_ops.stack(dim_0_sizes),
flat_ids.dtype))
ids_per_partition = num_total_ids // np
extras = num_total_ids % np
p_assignments = math_ops.maximum(
flat_ids // (ids_per_partition + 1),
(flat_ids - extras) // ids_per_partition)
# Emulate a conditional using a boolean indicator tensor
new_ids = array_ops.where(
p_assignments < extras, flat_ids % (ids_per_partition + 1),
(flat_ids - extras) % ids_per_partition)
else:
raise ValueError(
f"Unrecognized partition strategy: {partition_strategy}."
"Must be one of either `mod` or `div`.")
# Cast partition assignments to int32 for use in dynamic_partition.
# There really should not be more than 2^32 partitions.
p_assignments = math_ops.cast(p_assignments, dtypes.int32)
# Partition list of ids based on assignments into np separate lists
gather_ids = data_flow_ops.dynamic_partition(
new_ids, p_assignments, np)
# Similarly, partition the original indices.
pindices = data_flow_ops.dynamic_partition(original_indices,
p_assignments, np)
# Do np separate lookups, finding embeddings for plist[p] in params[p]
partitioned_result = []
for p in range(np):
pids = gather_ids[p]
with ops.device_v2(None):
with _colocate_with(params[p]):
result = array_ops.gather(params[p], pids)
if transform_fn:
# If transform_fn is provided, the clip_by_norm precedes
# the transform and hence must be co-located. See below
# for the counterpart if transform_fn is not provided.
result = transform_fn(_clip(
result, pids, max_norm))
partitioned_result.append(result)
# Stitch these back together
ret = data_flow_ops.parallel_dynamic_stitch(pindices,
partitioned_result,
name=name)
# Determine the static element shape.
if transform_fn is None:
element_shape_s = params[0].get_shape()[1:]
for p in params[1:]:
element_shape_s = element_shape_s.merge_with(
p.get_shape()[1:])
else:
element_shape_s = ret.get_shape()[1:]
# Compute the dynamic element shape.
if element_shape_s.is_fully_defined():
element_shape_d = element_shape_s
elif transform_fn is None:
# It's important that we compute params[0].shape on the right device
# to avoid data motion.
with _colocate_with(params[0]):
params_shape = array_ops.shape(params[0])
element_shape_d = params_shape[1:]
else:
element_shape_d = array_ops.shape(ret)[1:]
# Reshape to reverse the flattening of ids.
ret = array_ops.reshape(
ret,
array_ops.concat([array_ops.shape(ids), element_shape_d], 0))
# Normally the reshape is sufficient, but setting shape explicitly
# teaches shape inference that params[1:].get_shape() matters
# (in the case that transform_fn is None).
ret.set_shape(ids.get_shape().concatenate(element_shape_s))
if not transform_fn:
# If transform_fn was provided, the clip_by_norm was done above.
ret = _clip(ret, ids, max_norm)
return ret
def _interpolate_bilinear(grid,
query_points,
name='interpolate_bilinear',
indexing='ij'):
"""Similar to Matlab's interp2 function.
Finds values for query points on a grid using bilinear interpolation.
Args:
grid: a 4-D float `Tensor` of shape `[batch, height, width, channels]`.
query_points: a 3-D float `Tensor` of N points with shape `[batch, N, 2]`.
name: a name for the operation (optional).
indexing: whether the query points are specified as row and column (ij),
or Cartesian coordinates (xy).
Returns:
values: a 3-D `Tensor` with shape `[batch, N, channels]`
Raises:
ValueError: if the indexing mode is invalid, or if the shape of the inputs
invalid.
"""
if indexing != 'ij' and indexing != 'xy':
raise ValueError('Indexing mode must be \'ij\' or \'xy\'')
with ops.name_scope(name):
grid = ops.convert_to_tensor(grid)
query_points = ops.convert_to_tensor(query_points)
shape = grid.get_shape().as_list()
if len(shape) != 4:
msg = 'Grid must be 4 dimensional. Received size: '
raise ValueError(msg + str(grid.get_shape()))
batch_size, height, width, channels = shape
query_type = query_points.dtype
grid_type = grid.dtype
if (len(query_points.get_shape()) != 3
or query_points.get_shape()[2].value != 2):
msg = (
'Query points must be 3 dimensional and size 2 in dim 2. Received '
'size: ')
raise ValueError(msg + str(query_points.get_shape()))
_, num_queries, _ = query_points.get_shape().as_list()
if height < 2 or width < 2:
msg = 'Grid must be at least batch_size x 2 x 2 in size. Received size: '
raise ValueError(msg + str(grid.get_shape()))
alphas = []
floors = []
ceils = []
index_order = [0, 1] if indexing == 'ij' else [1, 0]
unstacked_query_points = array_ops.unstack(query_points, axis=2)
for dim in index_order:
with ops.name_scope('dim-' + str(dim)):
queries = unstacked_query_points[dim]
size_in_indexing_dimension = shape[dim + 1]
# max_floor is size_in_indexing_dimension - 2 so that max_floor + 1
# is still a valid index into the grid.
max_floor = math_ops.cast(size_in_indexing_dimension - 2,
query_type)
min_floor = constant_op.constant(0.0, dtype=query_type)
floor = math_ops.minimum(
math_ops.maximum(min_floor, math_ops.floor(queries)),
max_floor)
int_floor = math_ops.cast(floor, dtypes.int32)
floors.append(int_floor)
ceil = int_floor + 1
ceils.append(ceil)
# alpha has the same type as the grid, as we will directly use alpha
# when taking linear combinations of pixel values from the image.
alpha = math_ops.cast(queries - floor, grid_type)
min_alpha = constant_op.constant(0.0, dtype=grid_type)
max_alpha = constant_op.constant(1.0, dtype=grid_type)
alpha = math_ops.minimum(math_ops.maximum(min_alpha, alpha),
max_alpha)
# Expand alpha to [b, n, 1] so we can use broadcasting
# (since the alpha values don't depend on the channel).
alpha = array_ops.expand_dims(alpha, 2)
alphas.append(alpha)
if batch_size * height * width > np.iinfo(np.int32).max / 8:
error_msg = """The image size or batch size is sufficiently large
that the linearized addresses used by array_ops.gather
may exceed the int32 limit."""
raise ValueError(error_msg)
flattened_grid = array_ops.reshape(
grid, [batch_size * height * width, channels])
batch_offsets = array_ops.reshape(
math_ops.range(batch_size) * height * width, [batch_size, 1])
# This wraps array_ops.gather. We reshape the image data such that the
# batch, y, and x coordinates are pulled into the first dimension.
# Then we gather. Finally, we reshape the output back. It's possible this
# code would be made simpler by using array_ops.gather_nd.
def gather(y_coords, x_coords, name):
with ops.name_scope('gather-' + name):
linear_coordinates = batch_offsets + y_coords * width + x_coords
gathered_values = array_ops.gather(flattened_grid,
linear_coordinates)
return array_ops.reshape(gathered_values,
[batch_size, num_queries, channels])
# grab the pixel values in the 4 corners around each query point
top_left = gather(floors[0], floors[1], 'top_left')
top_right = gather(floors[0], ceils[1], 'top_right')
bottom_left = gather(ceils[0], floors[1], 'bottom_left')
bottom_right = gather(ceils[0], ceils[1], 'bottom_right')
# now, do the actual interpolation
with ops.name_scope('interpolate'):
interp_top = alphas[1] * (top_right - top_left) + top_left
interp_bottom = alphas[1] * (bottom_right -
bottom_left) + bottom_left
interp = alphas[0] * (interp_bottom - interp_top) + interp_top
return interp
def f(inputs):
num_steps, _ = inputs.shape[:2]
outputs = []
for t in math_ops.range(num_steps):
outputs.append(inputs[t])
return outputs
def _partial_shaped_bools():
rand_vect = math_ops.range(
random_ops.random_uniform(
shape=(), minval=2, maxval=3, dtype=dtypes.int32))
return array_ops.expand_dims_v2(rand_vect, 0) < 0
def _log_prob(self, x):
if self.cholesky_input_output_matrices:
x_sqrt = x
else:
# Complexity: O(nbk^3)
x_sqrt = linalg_ops.cholesky(x)
batch_shape = self.batch_shape_tensor()
event_shape = self.event_shape_tensor()
ndims = array_ops.rank(x_sqrt)
# sample_ndims = ndims - batch_ndims - event_ndims
sample_ndims = ndims - array_ops.shape(batch_shape)[0] - 2
sample_shape = array_ops.strided_slice(array_ops.shape(x_sqrt), [0],
[sample_ndims])
# We need to be able to pre-multiply each matrix by its corresponding
# batch scale matrix. Since a Distribution Tensor supports multiple
# samples per batch, this means we need to reshape the input matrix `x`
# so that the first b dimensions are batch dimensions and the last two
# are of shape [dimension, dimensions*number_of_samples]. Doing these
# gymnastics allows us to do a batch_solve.
#
# After we're done with sqrt_solve (the batch operation) we need to undo
# this reshaping so what we're left with is a Tensor partitionable by
# sample, batch, event dimensions.
# Complexity: O(nbk**2) since transpose must access every element.
scale_sqrt_inv_x_sqrt = x_sqrt
perm = array_ops.concat([
math_ops.range(sample_ndims, ndims),
math_ops.range(0, sample_ndims)
], 0)
scale_sqrt_inv_x_sqrt = array_ops.transpose(scale_sqrt_inv_x_sqrt,
perm)
shape = array_ops.concat(
(batch_shape,
(math_ops.cast(self.dimension, dtype=dtypes.int32), -1)), 0)
scale_sqrt_inv_x_sqrt = array_ops.reshape(scale_sqrt_inv_x_sqrt, shape)
# Complexity: O(nbM*k) where M is the complexity of the operator solving
# a vector system. E.g., for LinearOperatorDiag, each solve is O(k), so
# this complexity is O(nbk**2). For LinearOperatorLowerTriangular,
# each solve is O(k**2) so this step has complexity O(nbk^3).
scale_sqrt_inv_x_sqrt = self.scale_operator.solve(
scale_sqrt_inv_x_sqrt)
# Undo make batch-op ready.
# Complexity: O(nbk**2)
shape = array_ops.concat([batch_shape, event_shape, sample_shape], 0)
scale_sqrt_inv_x_sqrt = array_ops.reshape(scale_sqrt_inv_x_sqrt, shape)
perm = array_ops.concat([
math_ops.range(ndims - sample_ndims, ndims),
math_ops.range(0, ndims - sample_ndims)
], 0)
scale_sqrt_inv_x_sqrt = array_ops.transpose(scale_sqrt_inv_x_sqrt,
perm)
# Write V = SS', X = LL'. Then:
# tr[inv(V) X] = tr[inv(S)' inv(S) L L']
# = tr[inv(S) L L' inv(S)']
# = tr[(inv(S) L) (inv(S) L)']
# = sum_{ik} (inv(S) L)_{ik}**2
# The second equality follows from the cyclic permutation property.
# Complexity: O(nbk**2)
trace_scale_inv_x = math_ops.reduce_sum(
math_ops.square(scale_sqrt_inv_x_sqrt), axis=[-2, -1])
# Complexity: O(nbk)
half_log_det_x = math_ops.reduce_sum(math_ops.log(
array_ops.matrix_diag_part(x_sqrt)),
axis=[-1])
# Complexity: O(nbk**2)
log_prob = ((self.df - self.dimension - 1.) * half_log_det_x -
0.5 * trace_scale_inv_x - self.log_normalization())
# Set shape hints.
# Try to merge what we know from the input then what we know from the
# parameters of this distribution.
if x.get_shape().ndims is not None:
log_prob.set_shape(x.get_shape()[:-2])
if (log_prob.get_shape().ndims is not None
and self.batch_shape.ndims is not None
and self.batch_shape.ndims > 0):
log_prob.get_shape()[-self.batch_shape.ndims:].merge_with(
self.batch_shape)
return log_prob
def stack(self, name=None):
"""See TensorArray."""
with ops.colocate_with(self._handle):
with ops.name_scope(name, "TensorArrayStack", [self._handle]):
return self.gather(math_ops.range(0, self.size()), name=name)
def all_gather(self,
input_tensor,
axis,
communication_hint='AUTO',
timeout=0):
"""All-gather a dense tensor.
This method must be called inside a tf.function.
Args:
input_tensor: a dense tensor. It must have the same rank on all replicas,
and dimensions other than `axis` need to be the same as well.
axis: 0-D int32 Tensor. Dimension along which to gather. Must be in the
range [0, rank(value)).
communication_hint: string providing hint to runtime for choosing
collective implementation. Available options are `AUTO`, `NCCL`, and
`RING`.
timeout: a float. The timeout in seconds.
Returns:
The gathered Tensor.
Raises:
RuntimeError: if called in eager mode.
"""
if context.executing_eagerly():
raise RuntimeError('all_gather in eager mode is not supported')
instance_key_tensor = self._collective_keys.get_instance_key(
self._group_key, self._device)
instance_key_shape = self._collective_keys.get_instance_key(
self._group_key, self._device)
with ops.device(self._device):
# 1. Transpose
# E.g. Given an input_tensor with shape [2,2,5,1] and axis to gather is 3,
# we use perm_pre=[3 0 1 2] to reshape it to [1,2,2,5], which
# brings the 3rd dim first; afterwards we use perm_after=[1,2,3,0] to
# place it back.
perm_pre = array_ops.concat(
([axis], math_ops.range(axis),
math_ops.range(axis + 1, array_ops.rank(input_tensor))),
axis=0)
input_tensor_t = array_ops.transpose(input_tensor, perm=perm_pre)
# 2. Pad
gathered_shape = collective_ops.all_gather(
array_ops.expand_dims_v2(array_ops.shape_v2(input_tensor_t),
axis=0),
self._group_size,
self._group_key,
instance_key_shape,
communication_hint,
timeout=timeout)
first_dims = gathered_shape[:, 0]
full_axis_dim = math_ops.reduce_max(first_dims)
padded_input_tensor = _pad_util(input_tensor_t, full_axis_dim)
# 3. Gather
gather_padded_out_tensor = collective_ops.all_gather(
padded_input_tensor,
self._group_size,
self._group_key,
instance_key_tensor,
communication_hint,
timeout=timeout)
# 4. Unpad
split_tensors = []
for i in range(first_dims.shape[0]):
start_pos = i * full_axis_dim
split_tensors.append(
gather_padded_out_tensor[start_pos:start_pos +
first_dims[i]])
out_tensor_t = array_ops.concat(split_tensors, 0)
# 5. Transpose back
perm_after = array_ops.concat(
(math_ops.range(1, axis + 1), [0],
math_ops.range(axis + 1, array_ops.rank(input_tensor_t))),
axis=0)
return array_ops.transpose(out_tensor_t, perm=perm_after)
def train_steps(iterator):
for _ in math_ops.range(10):
train_step(iterator)
def test_sequential_model_save_and_load_weights(self):
if not context.executing_eagerly():
self.skipTest('Only test in eager mode.')
@tf.function
def slot_map_fn(x):
return tf.math.floormod(x, 2)
init = tf.keras.initializers.RandomNormal(seed=0)
model = get_sequential_model(de.keras.layers.FieldWiseEmbedding,
4,
2,
slot_map_fn,
bp_v2=False,
initializer=init,
name='pc053')
optmz = tf.keras.optimizers.Adam(learning_rate=1E-2, amsgrad=True)
optmz = de.DynamicEmbeddingOptimizer(optmz)
emb_layer = model.layers[0]
model.compile(optimizer=optmz, loss='binary_crossentropy')
start = 0
batch_size = 10
for i in range(1, 10):
x = math_ops.range(start,
start + batch_size * i,
dtype=dtypes.int64)
x = tf.reshape(x, (batch_size, -1))
start += batch_size * i
y = tf.zeros((batch_size, 1), dtype=dtypes.float32)
model.fit(x, y, verbose=0)
ids = tf.range(0, 10, dtype=tf.int64)
ids = tf.reshape(ids, (1, -1))
expected = model(ids)
save_dir = tempfile.mkdtemp(prefix='/tmp/')
options = tf.saved_model.SaveOptions(namespace_whitelist=['TFRA'])
model.save(save_dir, signatures=None, options=options)
copied_init = tf.keras.initializers.RandomNormal(seed=0)
new_model = get_sequential_model(de.keras.layers.FieldWiseEmbedding,
4,
2,
slot_map_fn,
bp_v2=False,
initializer=copied_init,
name='pc053')
new_emb_layer = new_model.layers[0]
new_model.load_weights(save_dir)
evaluate = new_model(ids)
self.assertAllEqual(evaluate, expected)
keys, values = emb_layer.params.export()
seq = tf.argsort(keys)
keys = tf.sort(keys)
values = tf.gather(values, seq)
new_keys, new_values = new_emb_layer.params.export()
seq = tf.argsort(new_keys)
new_keys = tf.sort(new_keys)
new_values = tf.gather(new_values, seq)
self.assertAllEqual(keys, new_keys)
self.assertAllEqual(values, new_values)
def random_deformation_momentum(shape, std, distance, stepsize=0.1):
r"""Create a random diffeomorphic deformation.
Parameters
----------
shape : sequence of 3 ints
Batch, height and width.
std : float
Correlation distance for the linear deformations.
distance : float
Expected total effective distance for the deformation.
stepsize : float
How large each step should be (as a propotion of ``std``).
Returns:
The end generated deformation field (Batch, height, width)
Notes
-----
``distance`` should typically not be more than a small fraction of the
sidelength of the image.
The computational time is is propotional to
.. math::
\frac{distance}{std * stepsize}
"""
grid_x, grid_y = array_ops.meshgrid(math_ops.range(shape[2]),
math_ops.range(shape[1]))
grid_x = tf.cast(grid_x[None, ..., None], 'float32')
grid_y = tf.cast(grid_y[None, ..., None], 'float32')
base_coordinates = tf.concat([grid_y, grid_x], axis=-1)
base_coordinates = tf.repeat(base_coordinates, shape[0], axis=0)
coordinates = tf.identity(base_coordinates)
# Create mask to stop movement at edges
mask = (tf.cos(
(grid_x - shape[2] / 2 + 1) * np.pi / (shape[2] + 2)) * tf.cos(
(grid_y - shape[1] / 2 + 1) * np.pi / (shape[1] + 2)))**(0.25)
# Total distance is given by std * n_steps * dt, we use this
# to work out the exact numbers.
n_steps = tf.cast(tf.math.ceil(distance / (std * stepsize)), 'int32')
dt = distance / (tf.cast(n_steps, 'float32') * std)
# Scale to get std 1 after smoothing
C = np.sqrt(2 * np.pi) * std**2
# Multiply by dt here to keep values small-ish for numerical purposes
momenta = dt * C * tf.random.normal(shape=[*shape, 2])
# Using a while loop, generate the deformation step-by-step.
def cond(i, from_coordinates, momenta):
return i < n_steps
def body(i, from_coordinates, momenta):
v = mask * gausssmooth(momenta, std)
d1 = matmul_transposed(jacobian(momenta), v)
d2 = matmul(jacobian(v), momenta)
d3 = div(v) * momenta
momenta = momenta - dt * (d1 + d2 + d3)
v = dense_image_warp(v, from_coordinates - base_coordinates)
from_coordinates = dense_image_warp(from_coordinates, v)
return i + 1, from_coordinates, momenta
i = tf.constant(0, dtype=tf.int32)
i, from_coordinates, momenta = tf.while_loop(cond, body,
[i, coordinates, momenta])
from_total_offset = from_coordinates - base_coordinates
return from_total_offset
def embedding_lookup(
params,
ids,
partition_strategy="mod",
name=None,
validate_indices=True, # pylint: disable=unused-argument
max_norm=None):
"""Looks up `ids` in a list of embedding tensors.
This function is used to perform parallel lookups on the list of
tensors in `params`. It is a generalization of
@{tf.gather}, where `params` is
interpreted as a partitioning of a large embedding tensor. `params` may be
a `PartitionedVariable` as returned by using `tf.get_variable()` with a
partitioner.
If `len(params) > 1`, each element `id` of `ids` is partitioned between
the elements of `params` according to the `partition_strategy`.
In all strategies, if the id space does not evenly divide the number of
partitions, each of the first `(max_id + 1) % len(params)` partitions will
be assigned one more id.
If `partition_strategy` is `"mod"`, we assign each id to partition
`p = id % len(params)`. For instance,
13 ids are split across 5 partitions as:
`[[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]]`
If `partition_strategy` is `"div"`, we assign ids to partitions in a
contiguous manner. In this case, 13 ids are split across 5 partitions as:
`[[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]]`
The results of the lookup are concatenated into a dense
tensor. The returned tensor has shape `shape(ids) + shape(params)[1:]`.
Args:
params: A single tensor representing the complete embedding tensor,
or a list of P tensors all of same shape except for the first dimension,
representing sharded embedding tensors. Alternatively, a
`PartitionedVariable`, created by partitioning along dimension 0. Each
element must be appropriately sized for the given `partition_strategy`.
ids: A `Tensor` with type `int32` or `int64` containing the ids to be looked
up in `params`.
partition_strategy: A string specifying the partitioning strategy, relevant
if `len(params) > 1`. Currently `"div"` and `"mod"` are supported. Default
is `"mod"`.
name: A name for the operation (optional).
validate_indices: DEPRECATED. If this operation is assigned to CPU, values
in `indices` are always validated to be within range. If assigned to GPU,
out-of-bound indices result in safe but unspecified behavior, which may
include raising an error.
max_norm: If not None, embedding values are l2-normalized to the value of
max_norm.
Returns:
A `Tensor` with the same type as the tensors in `params`.
Raises:
ValueError: If `params` is empty.
"""
if params is None or params in ((), []):
raise ValueError("Need at least one param")
if isinstance(params, variables.PartitionedVariable):
params = list(params) # Iterate to get the underlying Variables.
if not isinstance(params, list):
params = [params]
def maybe_normalize(x):
if max_norm is not None:
if x.get_shape().ndims is not None:
ndims = x.get_shape().ndims
else:
ndims = array_ops.size(array_ops.shape(x))
return clip_ops.clip_by_norm(x,
max_norm,
axes=list(range(1, ndims)))
return x
with ops.name_scope(name, "embedding_lookup", params + [ids]) as name:
np = len(params) # Number of partitions
# Preserve the resource variable status to avoid accidental dense reads.
if not any(
isinstance(p, resource_variable_ops.ResourceVariable)
for p in params):
params = ops.convert_n_to_tensor_or_indexed_slices(params,
name="params")
if np == 1:
with ops.colocate_with(params[0]):
return maybe_normalize(_do_gather(params[0], ids, name=name))
else:
ids = ops.convert_to_tensor(ids, name="ids")
flat_ids = array_ops.reshape(ids, [-1])
original_indices = math_ops.range(array_ops.size(flat_ids))
# Create p_assignments and set new_ids depending on the strategy.
if partition_strategy == "mod":
p_assignments = flat_ids % np
new_ids = flat_ids // np
elif partition_strategy == "div":
# Compute num_total_ids as the sum of dim-0 of params, then assign to
# partitions based on a constant number of ids per partition. Optimize
# if we already know the full shape statically.
dim_0_size = params[0].get_shape()[0]
for p in xrange(1, np):
dim_0_size += params[p].get_shape()[0]
if dim_0_size.value:
num_total_ids = constant_op.constant(
dim_0_size.value, flat_ids.dtype)
else:
dim_0_sizes = []
for p in xrange(np):
if params[p].get_shape()[0].value is not None:
dim_0_sizes.append(params[p].get_shape()[0].value)
else:
with ops.colocate_with(params[p]):
dim_0_sizes.append(
array_ops.shape(params[p])[0])
num_total_ids = math_ops.reduce_sum(
math_ops.cast(array_ops.stack(dim_0_sizes),
flat_ids.dtype))
ids_per_partition = num_total_ids // np
extras = num_total_ids % np
p_assignments = math_ops.maximum(
flat_ids // (ids_per_partition + 1),
(flat_ids - extras) // ids_per_partition)
# Emulate a conditional using a boolean indicator tensor
is_in_first_extras_partitions = math_ops.cast(
p_assignments < extras, flat_ids.dtype)
new_ids = (is_in_first_extras_partitions *
(flat_ids % (ids_per_partition + 1)) +
(1 - is_in_first_extras_partitions) *
((flat_ids - extras) % ids_per_partition))
else:
raise ValueError("Unrecognized partition strategy: " +
partition_strategy)
# Cast partition assignments to int32 for use in dynamic_partition.
# There really should not be more than 2^32 partitions.
p_assignments = math_ops.cast(p_assignments, dtypes.int32)
# Partition list of ids based on assignments into np separate lists
gather_ids = data_flow_ops.dynamic_partition(
new_ids, p_assignments, np)
# Similarly, partition the original indices.
pindices = data_flow_ops.dynamic_partition(original_indices,
p_assignments, np)
# Do np separate lookups, finding embeddings for plist[p] in params[p]
partitioned_result = []
for p in xrange(np):
with ops.colocate_with(params[p]):
partitioned_result.append(
_do_gather(params[p], gather_ids[p]))
# Stitch these back together
ret = data_flow_ops.dynamic_stitch(pindices,
partitioned_result,
name=name)
# Reshape to reverse the flattening of ids.
element_shape = params[0].get_shape()[1:]
for p in params[1:]:
element_shape = element_shape.merge_with(p.get_shape()[1:])
if element_shape.is_fully_defined():
ret = array_ops.reshape(
ret,
array_ops.concat([array_ops.shape(ids), element_shape], 0))
else:
# It's important that we compute params[0].shape on the right device
# to avoid data motion.
with ops.colocate_with(params[0]):
params_shape = array_ops.shape(params[0])
ret = array_ops.reshape(
ret,
array_ops.concat([
array_ops.shape(ids),
array_ops.slice(params_shape, [1], [-1])
], 0))
# output shape = ids.shape + params[*].shape[1:]
# Normally the reshape is sufficient, but setting shape explicitly
# teaches shape inference that params[1:].get_shape() matters.
ret.set_shape(ids.get_shape().concatenate(element_shape))
return maybe_normalize(ret)
def __init__(self,
distribution,
bijector=None,
batch_shape=None,
event_shape=None,
validate_args=False,
name=None):
"""Construct a Transformed Distribution.
Args:
distribution: The base distribution instance to transform. Typically an
instance of `Distribution`.
bijector: The object responsible for calculating the transformation.
Typically an instance of `Bijector`. `None` means `Identity()`.
batch_shape: `integer` vector `Tensor` which overrides `distribution`
`batch_shape`; valid only if `distribution.is_scalar_batch()`.
event_shape: `integer` vector `Tensor` which overrides `distribution`
`event_shape`; valid only if `distribution.is_scalar_event()`.
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.
name: Python `str` name prefixed to Ops created by this class. Default:
`bijector.name + distribution.name`.
"""
parameters = dict(locals())
name = name or (("" if bijector is None else bijector.name) +
distribution.name)
with ops.name_scope(name, values=[event_shape, batch_shape]) as name:
# For convenience we define some handy constants.
self._zero = constant_op.constant(0,
dtype=dtypes.int32,
name="zero")
self._empty = constant_op.constant([],
dtype=dtypes.int32,
name="empty")
if bijector is None:
bijector = identity_bijector.Identity(
validate_args=validate_args)
# We will keep track of a static and dynamic version of
# self._is_{batch,event}_override. This way we can do more prior to graph
# execution, including possibly raising Python exceptions.
self._override_batch_shape = self._maybe_validate_shape_override(
batch_shape, distribution.is_scalar_batch(), validate_args,
"batch_shape")
self._is_batch_override = _logical_not(
_logical_equal(_ndims_from_shape(self._override_batch_shape),
self._zero))
self._is_maybe_batch_override = bool(
tensor_util.constant_value(self._override_batch_shape) is None
or tensor_util.constant_value(
self._override_batch_shape).size != 0)
self._override_event_shape = self._maybe_validate_shape_override(
event_shape, distribution.is_scalar_event(), validate_args,
"event_shape")
self._is_event_override = _logical_not(
_logical_equal(_ndims_from_shape(self._override_event_shape),
self._zero))
self._is_maybe_event_override = bool(
tensor_util.constant_value(self._override_event_shape) is None
or tensor_util.constant_value(
self._override_event_shape).size != 0)
# To convert a scalar distribution into a multivariate distribution we
# will draw dims from the sample dims, which are otherwise iid. This is
# easy to do except in the case that the base distribution has batch dims
# and we're overriding event shape. When that case happens the event dims
# will incorrectly be to the left of the batch dims. In this case we'll
# cyclically permute left the new dims.
self._needs_rotation = _logical_and(
self._is_event_override, _logical_not(self._is_batch_override),
_logical_not(distribution.is_scalar_batch()))
override_event_ndims = _ndims_from_shape(
self._override_event_shape)
self._rotate_ndims = _pick_scalar_condition(
self._needs_rotation, override_event_ndims, 0)
# We'll be reducing the head dims (if at all), i.e., this will be []
# if we don't need to reduce.
self._reduce_event_indices = math_ops.range(
self._rotate_ndims - override_event_ndims, self._rotate_ndims)
self._distribution = distribution
self._bijector = bijector
super(TransformedDistribution, self).__init__(
dtype=self._distribution.dtype,
reparameterization_type=self._distribution.reparameterization_type,
validate_args=validate_args,
allow_nan_stats=self._distribution.allow_nan_stats,
parameters=parameters,
# We let TransformedDistribution access _graph_parents since this class
# is more like a baseclass than derived.
graph_parents=(
distribution._graph_parents + # pylint: disable=protected-access
bijector.graph_parents),
name=name)
def sequence_loss_by_example(logits,
targets,
weights,
num_decoder_symbols,
average_across_timesteps=True,
softmax_loss_function=None,
name=None):
"""Weighted cross-entropy loss for a sequence of logits (per example).
Args:
logits: list of 2D Tensors of shape [batch_size x num_decoder_symbols].
targets: list of 1D batch-sized int32 Tensors of the same length as logits.
weights: list of 1D batch-sized float-Tensors of the same length as logits.
num_decoder_symbols: integer, number of decoder symbols (output classes).
average_across_timesteps: If set, divide the returned cost by the total
label weight.
softmax_loss_function: function (inputs-batch, labels-batch) -> loss-batch
to be used instead of the standard softmax (the default if this is None).
name: optional name for this operation, default: "sequence_loss_by_example".
Returns:
1D batch-sized float Tensor: the log-perplexity for each sequence.
Raises:
ValueError: if len(logits) is different from len(targets) or len(weights).
"""
if len(targets) != len(logits) or len(weights) != len(logits):
raise ValueError(
"Lengths of logits, weights, and targets must be the same "
"%d, %d, %d." % (len(logits), len(weights), len(targets)))
with ops.op_scope(logits + targets + weights, name,
"sequence_loss_by_example"):
batch_size = array_ops.shape(targets[0])[0]
log_perp_list = []
length = batch_size * num_decoder_symbols
for i in xrange(len(logits)):
if softmax_loss_function is None:
# TODO(lukaszkaiser): There is no SparseCrossEntropy in TensorFlow, so
# we need to first cast targets into a dense representation, and as
# SparseToDense does not accept batched inputs, we need to do this by
# re-indexing and re-sizing. When TensorFlow adds SparseCrossEntropy,
# rewrite this method.
indices = targets[i] + num_decoder_symbols * math_ops.range(
batch_size)
with ops.device(
"/cpu:0"): # Sparse-to-dense must be on CPU for now.
dense = sparse_ops.sparse_to_dense(
indices, array_ops.expand_dims(length, 0), 1.0, 0.0)
target = array_ops.reshape(dense, [-1, num_decoder_symbols])
crossent = nn_ops.softmax_cross_entropy_with_logits(
logits[i],
target,
name="SequenceLoss/CrossEntropy{0}".format(i))
else:
crossent = softmax_loss_function(logits[i], targets[i])
log_perp_list.append(crossent * weights[i])
log_perps = math_ops.add_n(log_perp_list)
if average_across_timesteps:
total_size = math_ops.add_n(weights)
total_size += 1e-12 # Just to avoid division by 0 for all-0 weights.
log_perps /= total_size
return log_perps
def connected_components(images):
"""Labels the connected components in a batch of images.
A component is a set of pixels in a single input image, which are all adjacent
and all have the same non-zero value. The components using a squared
connectivity of one (all True entries are joined with their neighbors above,
below, left, and right). Components across all images have consecutive ids 1
through n. Components are labeled according to the first pixel of the
component appearing in row-major order (lexicographic order by
image_index_in_batch, row, col). Zero entries all have an output id of 0.
This op is equivalent with `scipy.ndimage.measurements.label` on a 2D array
with the default structuring element (which is the connectivity used here).
Args:
images: A 2D (H, W) or 3D (N, H, W) Tensor of boolean image(s).
Returns:
Components with the same shape as `images`. False entries in `images` have
value 0, and all True entries map to a component id > 0.
Raises:
TypeError: if `images` is not 2D or 3D.
"""
with ops.name_scope("connected_components"):
image_or_images = ops.convert_to_tensor(images, name="images")
if len(image_or_images.get_shape()) == 2:
images = image_or_images[None, :, :]
elif len(image_or_images.get_shape()) == 3:
images = image_or_images
else:
raise TypeError(
"images should have rank 2 (HW) or 3 (NHW). Static shape is %s"
% image_or_images.get_shape())
components = gen_image_ops.image_connected_components(images)
# TODO(ringwalt): Component id renaming should be done in the op, to avoid
# constructing multiple additional large tensors.
components_flat = array_ops.reshape(components, [-1])
unique_ids, id_index = array_ops.unique(components_flat)
id_is_zero = array_ops.where(math_ops.equal(unique_ids, 0))[:, 0]
# Map each nonzero id to consecutive values.
nonzero_consecutive_ids = math_ops.range(
array_ops.shape(unique_ids)[0] -
array_ops.shape(id_is_zero)[0]) + 1
def no_zero():
# No need to insert a zero into the ids.
return nonzero_consecutive_ids
def has_zero():
# Insert a zero in the consecutive ids where zero appears in unique_ids.
# id_is_zero has length 1.
zero_id_ind = math_ops.to_int32(id_is_zero[0])
ids_before = nonzero_consecutive_ids[:zero_id_ind]
ids_after = nonzero_consecutive_ids[zero_id_ind:]
return array_ops.concat([ids_before, [0], ids_after], axis=0)
new_ids = control_flow_ops.cond(
math_ops.equal(array_ops.shape(id_is_zero)[0], 0), no_zero,
has_zero)
components = array_ops.reshape(array_ops.gather(new_ids, id_index),
array_ops.shape(components))
if len(image_or_images.get_shape()) == 2:
return components[0, :, :]
else:
return components
def test_model_save_and_load(self):
if not context.executing_eagerly():
self.skipTest('Only test in eager mode.')
@tf.function
def slot_map_fn(x):
return tf.math.floormod(x, 2)
init = tf.keras.initializers.RandomNormal(seed=0)
class MyModel(tf.keras.Model):
def __init__(self):
super(MyModel, self).__init__()
self.l0 = tf.keras.layers.InputLayer(input_shape=(None, ),
dtype=tf.int64)
self.l1 = de.keras.layers.FieldWiseEmbedding(4,
2,
slot_map_fn,
bp_v2=False,
initializer=init,
name='sl337')
self.l2 = tf.keras.layers.Flatten()
self.l3 = tf.keras.layers.Dense(32, 'relu')
self.l4 = tf.keras.layers.Dense(1, 'sigmoid')
def call(self, x):
return self.l4(self.l3(self.l2(self.l1(self.l0(x)))))
model = MyModel()
optmz = tf.keras.optimizers.Adam(1E-3)
optmz = de.DynamicEmbeddingOptimizer(optmz)
model.compile(optimizer=optmz, loss='binary_crossentropy')
start = 0
batch_size = 10
for i in range(1, 10):
x = math_ops.range(start,
start + batch_size * i,
dtype=dtypes.int64)
x = tf.reshape(x, (batch_size, -1))
start += batch_size * i
y = tf.zeros((batch_size, 1), dtype=dtypes.float32)
model.fit(x, y, verbose=0)
ids = tf.range(0, 10, dtype=tf.int64)
ids = tf.reshape(ids, (2, -1))
expected = model(ids)
save_dir = tempfile.mkdtemp(prefix='/tmp/')
options = tf.saved_model.SaveOptions(namespace_whitelist=['TFRA'])
@tf.function(
input_signature=[tf.TensorSpec((None, None), dtype=tf.int64)])
def foo(x):
return model(x)
model.save(save_dir, signatures=foo, options=options)
new_model = tf.saved_model.load(save_dir)
sig = new_model.signatures['serving_default']
evaluated = sig(ids)['output_0']
self.assertAllClose(expected, evaluated, 1E-7, 1E-7)
def _Range(self, start, limit, delta):
with self.test_session(use_gpu=True):
tf_ans = math_ops.range(start, limit, delta, name="range")
self.assertEqual([len(np.arange(start, limit, delta))],
tf_ans.get_shape())
return tf_ans.eval()
def testLimitOnly(self):
with self.test_session(use_gpu=True):
self.assertAllEqual(np.arange(5), math_ops.range(5).eval())