def testFoldl_Simple(self):
with self.test_session():
elems = tf.constant([1, 2, 3, 4, 5, 6], name="data")
r = tf.foldl(lambda a, x: tf.mul(tf.add(a, x), 2), elems)
self.assertAllEqual(208, r.eval())
r = tf.foldl(lambda a, x: tf.mul(tf.add(a, x), 2), elems, initializer=10)
self.assertAllEqual(880, r.eval())
def testFoldl_Simple(self):
with self.test_session():
elems = tf.constant([1, 2, 3, 4, 5, 6], name="data")
r = tf.foldl(lambda a, x: tf.mul(tf.add(a, x), 2), elems)
self.assertAllEqual(208, r.eval())
r = tf.foldl(
lambda a, x: tf.mul(tf.add(a, x), 2), elems, initializer=10)
self.assertAllEqual(880, r.eval())
def __init__(self, window=15, **kwargs):
super(SegmentInputLayer, self).__init__(**kwargs)
sizes = self.input_data.size_placeholder[0]
new_sizes = batch_sizes_after_windowing(sizes, window)
def fold_data(acc, x):
batch_idx = x[0]
num_frames = x[1]
res = tf.expand_dims(tf.range(num_frames), -1) # start times
res = tf.tile(res, [1, window]) # fill add time dimension
res += tf.range(window) # add offsets
res = tf.where(res >= num_frames, tf.zeros_like(res), res) # filter frames that go past the end
if self.input_data.is_batch_major:
res = tf.stack([tf.ones_like(res) * batch_idx, res], axis=2) # add batch_index in first dim
else:
res = tf.stack([res, tf.ones_like(res) * batch_idx], axis=2) # add batch_index in second dim
return tf.concat([acc, res], 0)
initial = tf.placeholder_with_default((tf.zeros([0, window, 2], dtype=tf.int32)), [None, window, 2])
indices = tf.foldl(fold_data, tf.stack([tf.range(tf.shape(sizes)[0]), sizes], axis=1), initial, name='fold_data')
if self.input_data.is_time_major:
indices = tf.transpose(indices, [1, 0, 2])
self.output.placeholder = tf.gather_nd(self.input_data.placeholder, indices)
self.output.size_placeholder[0] = new_sizes
def call(self,
decoder_input,
encoder_output,
decoder_self_attention_bias,
training=False,
mask=None):
if self._hparams.recurrence_type == "basic":
x = decoder_input
ut_initializer = (x, x, x) # (state, input, memory)
ut_decoder_unit = functools.partial(
universal_transformer_basic,
encoder_output=encoder_output,
decoder_self_attention_bias=decoder_self_attention_bias,
step_preprocess=self.step_preprocess,
attention_unit=self.transformer_decoder_attention_unit,
ffn_unit=self.transformer_decoder_ffn_unit,
training=training,
)
output, _, extra_output = tf.foldl(
ut_decoder_unit, tf.range(self._hparams.num_rec_steps),
initializer=ut_initializer,
)
output = self.decoder_normalizer(None, output, training)
return output, extra_output
if self._hparams.recurrence_type == "act":
output, extra_output = universal_transformer_act(decoder_input, encoder_output,
decoder_self_attention_bias, self.step_preprocess,
self.transformer_decoder_attention_unit,
self.transformer_decoder_ffn_unit, self.halting_unit,
self._hparams, training)
output = self.decoder_normalizer(None, output, training)
return output, extra_output
def average_histogram(data, num_bins, num_shifts):
# Set up the range tensor
data_min = tf.reduce_min(data)
data_max = tf.reduce_max(data)
data_range = [data_min, data_max]
range_tensor_2d = tf.expand_dims(data_range, 0)
# Set up the shift tensor
half_bin = (data_max - data_min) / (num_bins * 2)
shift = tf.linspace(-half_bin, half_bin, num_shifts)
shift_tensor_2d = tf.expand_dims(shift, 1)
# Create the range elements for the shifts using broadcast addition
range_elements = tf.add(range_tensor_2d, shift_tensor_2d)
# Performs one histogram with the given range, accumulating the counts
def one_histogram(accumulator, element):
hist = tf.histogram_fixed_width(data, element, num_bins)
return tf.add(accumulator, hist)
# Perform the shifted histograms and divide by the total count
initial_accumulator = tf.zeros([num_bins], tf.int32)
sum_hist = tf.foldl(one_histogram, range_elements, initial_accumulator)
total = tf.to_float(tf.reduce_sum(sum_hist))
avg_hist = tf.div(tf.to_float(sum_hist), total)
return avg_hist
def stacked_conv(sparse_neighbors,
x_in,
x_out,
features_in,
bias,
kernels,
impl=None):
if not isinstance(kernels, (list, tuple)):
raise ValueError(
'kernels should be a list or tuple, got {}'.format(kernels))
N_out, _ = tf.unstack(sparse_neighbors.dense_shape)
_, F_out = bias.shape
x_in_stack, x_out_stack, kernel_stack, = _get_polynomial_stacks(
x_in, x_out, bias, kernels)
# replace `None`s with ones
in0 = tf.ones_like(x_in)
out0 = tf.ones_like(x_out)
x_in_stack = [in0 if x is None else x for x in x_in_stack]
x_out_stack = [out0 if x is None else x for x in x_out_stack]
x_in = tf.concat(x_in_stack, axis=0)
x_out = tf.concat(x_out_stack, axis=0)
kernel = tf.concat(kernel_stack, axis=0)
init = tf.zeros(shape=(N_out, F_out), dtype=features_in.dtype)
return tf.foldl(
lambda acc, args: inner(sparse_neighbors,
args[0],
args[1],
features_in,
args[2],
impl=impl), (x_in, x_out, kernel), init)
def __init__(self, window=15, **kwargs):
super(SegmentInputLayer, self).__init__(**kwargs)
sizes = self.input_data.size_placeholder[0]
new_sizes = batch_sizes_after_windowing(sizes, window)
def fold_data(acc, x):
batch_idx = x[0]
num_frames = x[1]
res = tf.expand_dims(tf.range(num_frames), -1) # start times
res = tf.tile(res, [1, window]) # fill add time dimension
res += tf.range(window) # add offsets
res = tf.where(res >= num_frames, tf.zeros_like(res),
res) # filter frames that go past the end
if self.input_data.is_batch_major:
res = tf.stack([tf.ones_like(res) * batch_idx, res],
axis=2) # add batch_index in first dim
else:
res = tf.stack([res, tf.ones_like(res) * batch_idx],
axis=2) # add batch_index in second dim
return tf.concat([acc, res], 0)
initial = tf_compat.v1.placeholder_with_default(
(tf.zeros([0, window, 2], dtype=tf.int32)), [None, window, 2])
indices = tf.foldl(fold_data,
tf.stack([tf.range(tf.shape(sizes)[0]), sizes],
axis=1),
initial,
name='fold_data')
if self.input_data.is_time_major:
indices = tf.transpose(indices, [1, 0, 2])
self.output.placeholder = tf.gather_nd(self.input_data.placeholder,
indices)
self.output.size_placeholder[0] = new_sizes
def naive_conv(sparse_neighbors, x_in, x_out, features_in, kernel, bias=None):
"""
Args:
sparse_neighbors: [N_out, N_in] sparse float tensor of weighted values.
Rows should sum to 1 (though this is not ef_inorced). Assumed to be
ordered (if not, use `tf.sparse.reorder(sparse_neighbors)`).
x_in: [D, N_in] float tensor of input cloud coordinates.
x_out: [D, N_out] float tensor of output cloud coordinates.
features_in: [N_in, F_in] float tensor of input features.
kernel: [D, F_in, F_out] float tensor of kernel.
Returns:
[N_out, F_out] float output features.
"""
D, F_in, F_out = _check_kernel_shape(x_in, features_in, kernel)
del D, F_in
i, j = tf.unstack(sparse_neighbors.indices, axis=-1)
dx = tf.gather(x_out, i, axis=1) - tf.gather(x_in, j, axis=1)
dx, kernel = _merge_kernel_and_bias(dx, kernel, bias)
del bias
dx = dx * tf.expand_dims(sparse_neighbors.values, axis=0)
gathered_features = tf.gather(features_in, j)
def get_kth_term(k, dx):
out = tf.matmul(gathered_features, k) * tf.expand_dims(dx, axis=-1)
reduced = tf.math.segment_sum(out, i)
return reduced
init = tf.zeros((sparse_neighbors.dense_shape[0], F_out),
dtype=features_in.dtype)
return tf.foldl(_sum_map(get_kth_term), (kernel, dx), init)
def top_kth_iterative(x, k):
"""Compute the k-th top element of x on the last axis iteratively.
This assumes values in x are non-negative, rescale if needed.
It is often faster than tf.nn.top_k for small k, especially if k < 30.
Note: this does not support back-propagation, it stops gradients!
Args:
x: a Tensor of non-negative numbers of type float.
k: a python integer.
Returns:
a float tensor of the same shape as x but with 1 on the last axis
that contains the k-th largest number in x.
"""
# The iterative computation is as follows:
#
# cur_x = x
# for _ in range(k):
# top_x = maximum of elements of cur_x on the last axis
# cur_x = cur_x where cur_x < top_x and 0 everywhere else (top elements)
#
# We encode this computation in a TF graph using tf.foldl, so the inner
# part of the above loop is called "next_x" and tf.foldl does the loop.
def next_x(cur_x, _):
top_x = tf.reduce_max(cur_x, axis=-1, keep_dims=True)
return cur_x * to_float(cur_x < top_x)
# We only do k-1 steps of the loop and compute the final max separately.
fin_x = tf.foldl(next_x,
tf.range(k - 1),
initializer=tf.stop_gradient(x),
parallel_iterations=2,
back_prop=False)
return tf.stop_gradient(tf.reduce_max(fin_x, axis=-1, keep_dims=True))
def tf_draw_masks_on_image(image,
mask,
color=None,
alpha=0.4,
name='summary_image_with_mask'):
if image.dtype is not tf.uint8:
image = tf.cast(image, tf.uint8)
if mask.dtype is not tf.uint8:
mask = tf.cast(mask, tf.uint8)
with tf.name_scope(name):
nr = tf.shape(mask)[0]
color_nr = tf.shape(color)[0]
indexs = tf.range(nr)
def fn(img, index):
m = mask[index]
c = color[index % color_nr]
def fn0():
return tf_draw_mask_on_image_array(img, m, c, alpha)
def fn1():
return img
return tf.cond(tf.greater(tf.reduce_max(m), 0), fn0, fn1)
image = tf.foldl(fn, elems=indexs, initializer=image, back_prop=False)
return image
def in_term_v2(sparse_neighbors, x_in, features_in, kernel):
"""
Implements x_in term of convolution.
term_{ip} = \\sum_{k,q,j} n_{ij} theta_{pqk} xin_{jk} f_{jq}
= \\sum_j n_{ij} \\sum_{k} xin_{jk} \\sum_{q} theta_{pqk} f_{jq}
Largest tensor: [N, D, F_out]. Good if F_out < F_in.
Args:
sparse_neighbors: [N_out, N_in] sparse float tensor of normalized
weighted values.
x_in: [D, N_in] float tensor of output cloud coordinates.
features_in: [N_in, F_in] float tensor of input features.
kernel: [D, F_in, F_out] float tensor corresponding to theta.
Returns:
[N_out, F_out] float output features.
"""
D, F_in, F_out = _check_kernel_shape(x_in, features_in, kernel)
del D
if F_in > F_out:
def get_kth_term(k, x):
return tf.matmul(features_in, k) * tf.expand_dims(x, axis=-1)
else:
def get_kth_term(k, x):
return tf.matmul(features_in * tf.expand_dims(x, axis=-1), k)
f = tf.foldl(
_sum_map(get_kth_term), (kernel, x_in),
tf.zeros((sparse_neighbors.dense_shape[1], F_out), features_in.dtype))
return tf.sparse.sparse_dense_matmul(sparse_neighbors, f)
def testFoldl():
x_h = tf.compat.v1.placeholder(tf.int32, [None, 5])
x = np.array([[3, 4, 0, 2, 1], [2, 4, 3, 0, 1]])
fd = tf.foldl(lambda a, b: tf.stack(a, tf.math.invert_permutation(b)), x_h)
with tf.compat.v1.Session() as sess:
r = sess.run(fd, feed_dict={x_h: x})
print(r)
def perform_filter_operation(Y, filter_matrix_conj, taps, delay):
"""
>>> D, T, taps, delay = 1, 10, 2, 1
>>> tf.enable_eager_execution()
>>> Y = tf.ones([D, T])
>>> filter_matrix_conj = tf.ones([taps, D, D])
>>> X = perform_filter_operation_v2(Y, filter_matrix_conj, taps, delay)
>>> X.shape
TensorShape([Dimension(1), Dimension(10)])
>>> X.numpy()
array([[ 1., 0., -1., -1., -1., -1., -1., -1., -1., -1.]], dtype=float32)
"""
dyn_shape = tf.shape(Y)
T = dyn_shape[1]
def add_tap(accumulated, tau_minus_delay):
new = tf.einsum(
'de,dt',
filter_matrix_conj[tau_minus_delay, :, :],
Y[:, :(T - delay - tau_minus_delay)]
)
paddings = tf.convert_to_tensor([[0, 0], [delay + tau_minus_delay, 0]])
new = tf.pad(new, paddings, "CONSTANT")
return accumulated + new
reverb_tail = tf.foldl(
add_tap, tf.range(0, taps),
initializer=tf.zeros_like(Y)
)
return Y - reverb_tail
def map_reduce_sum(map_fn,
inputs,
out_shape,
out_type,
parallel_iterations=None,
method="unstack"):
if method == "fold":
init = tf.zeros(out_shape, dtype=out_type)
return tf.foldl(lambda acc, args: acc + map_fn(args), inputs, init)
if method == "map":
return tf.reduce_sum(
tf.map_fn(map_fn,
inputs,
parallel_iterations=parallel_iterations,
dtype=out_type),
axis=0,
)
if method == "vmap":
return tf.reduce_sum(tf.vectorized_map(map_fn, inputs), axis=0)
if method == "unstack":
inputs = (tf.unstack(i, axis=0) for i in inputs)
return tf.add_n([map_fn(args) for args in zip(*inputs)])
options = ("fold", "map", "vmap", "unstack")
raise ValueError(f"Invalid method {method}: must be one of {options}")
def _foldl_repeat(values, repeats):
"""Repeat using `tf.foldl` with `TensorArray` concatenation.
Mostly just a fallback in case of issues importing `ragged_util.repeat`.
The ragged implementation is significantly faster for small-to-medium size
arrays, though benchmarking indicates this may be faster for generating arrays
of ~200k blocks of ~1k elements each or more.
"""
size = values.shape[0]
element_shape = (None,) + tuple(values.shape[1:])
acc0 = tf.TensorArray(
dtype=values.dtype,
size=size,
dynamic_size=False,
element_shape=element_shape,
infer_shape=False,
)
def fold_fn(out, args):
index, acc = out
repeat, value = args
value = tf.fill((repeat,), value)
acc = acc.write(index, value)
return index + 1, acc
acc = tf.foldl(
fold_fn,
(repeats, values),
initializer=(0, acc0),
parallel_iterations=8,
back_prop=False,
)[1]
return acc.concat()
def get_unique_elements_with_counts(
dataset: tf.data.Dataset,
max_string_length: Optional[int] = None) -> Tuple[tf.Tensor, tf.Tensor]:
"""Gets unique elements and their counts from the input `dataset`.
This method returns a tuple of `elements` and `counts`, where `elements` are
the unique elements in the dataset, and counts is the number of times each one
appears.
The input `dataset` must yield batched rank-1 tensors. This function reads
each coordinate of the tensor as an individual element and caps the total
number of elements to return.
Args:
dataset: A `tf.data.Dataset` to elements from. Element type must be
`tf.string`.
max_string_length: The maximum lenghth (in bytes) of strings in the dataset.
Strings longer than `max_string_length` will be truncated. Defaults to
`None`, which means there is no limit of the string length.
Returns:
elements: A rank-1 Tensor containing all the unique elements of the input
`dataset`.
counts: A rank-1 Tensor containing the counts for each of the elements in
`elements`.
Raises:
ValueError:
-- If the shape of elements in `dataset` is not rank 1
-- If `max_string_length` is not `None` and is less than 1.
TypeError: If `dataset.element_spec.dtype` must be `tf.string` is not
`tf.string`.
"""
if dataset.element_spec.shape.rank != 1:
raise ValueError('The shape of elements in `dataset` must be of rank 1, '
f' found rank = {dataset.element_spec.shape.rank}'
' instead.')
if max_string_length is not None and max_string_length < 1:
raise ValueError('`max_string_length` must be at least 1 when it is not'
' None.')
if dataset.element_spec.dtype != tf.string:
raise TypeError('`dataset.element_spec.dtype` must be `tf.string`, found'
f' element type {dataset.element_spec.dtype}')
all_elements = get_all_elements(dataset, max_string_length)
elements, indices = tf.unique(all_elements)
def accumulate_counts(counts, item):
item_one_hot = tf.one_hot(item, tf.shape(elements)[0], dtype=tf.int64)
return counts + item_one_hot
counts = tf.foldl(
accumulate_counts,
indices,
initializer=tf.zeros_like(elements, dtype=tf.int64))
return elements, counts
def test():
elems = tf.Variable(np.arange(10, dtype=np.float32))
elems = tf.identity(elems)
op_sum = tf.foldl(lambda a, x: a + x, elems)
op_mean = tf.reduce_mean(elems)
with tf.Session() as sess:
sess.run(tf.initialize_all_variables())
print sess.run([op_sum, op_mean])
def make_points_to_edges_indexs(self):
with tf.variable_scope("make_points_to_edges_index"):
datas = tf.ones(shape=[2, self.real_edge_nr], dtype=tf.int32)
datas = tf.foldl(self.make_point_to_edges_indexs,
elems=self.point_indexs,
initializer=datas,
back_prop=False)
return tf.unstack(datas, axis=0)
def testFoldl_Scoped(self):
with self.test_session() as sess:
with tf.variable_scope("root") as varscope:
elems = tf.constant([1, 2, 3, 4, 5, 6], name="data")
r = tf.foldl(simple_scoped_fn, elems)
# Check that we have the one variable we asked for here.
self.assertEqual(len(tf.trainable_variables()), 1)
self.assertEqual(tf.trainable_variables()[0].name, "root/body/two:0")
sess.run([tf.initialize_all_variables()])
self.assertAllEqual(208, r.eval())
# Now let's reuse our single variable.
varscope.reuse_variables()
r = tf.foldl(simple_scoped_fn, elems, initializer=10)
self.assertEqual(len(tf.trainable_variables()), 1)
self.assertAllEqual(880, r.eval())
def ph_blur(inputs, values_in, offsets, weights, nbs):
# TODO: we can parameterize all this !
def _blur_iter(prev, nbs):
n1 = tf.gather(prev, nbs[:, 0] + 1)
n2 = tf.gather(prev, nbs[:, 1] + 1)
return prev + 0.5 * tf.pad(n1 + n2, [[1, 1], [0, 0]])
return tf.foldl(_blur_iter, tf.transpose(nbs, [1, 0, 2]), values_in)
def testFoldShape(self):
with self.test_session():
x = tf.constant([[1, 2, 3], [4, 5, 6]])
def fn(_, current_input):
return current_input
initializer = tf.constant([0, 0, 0])
y = tf.foldl(fn, x, initializer=initializer)
self.assertAllEqual(y.get_shape(), y.eval().shape)
def testFoldl_Scoped(self):
with self.test_session() as sess:
with tf.variable_scope("root") as varscope:
elems = tf.constant([1, 2, 3, 4, 5, 6], name="data")
r = tf.foldl(simple_scoped_fn, elems)
# Check that we have the one variable we asked for here.
self.assertEqual(len(tf.trainable_variables()), 1)
self.assertEqual(tf.trainable_variables()[0].name, "root/body/two:0")
sess.run([tf.initialize_all_variables()])
self.assertAllEqual(208, r.eval())
# Now let's reuse our single variable.
varscope.reuse_variables()
r = tf.foldl(simple_scoped_fn, elems, initializer=10)
self.assertEqual(len(tf.trainable_variables()), 1)
self.assertAllEqual(880, r.eval())
def test():
elems = tf.Variable(np.arange(10, dtype=np.float32))
elems = tf.identity(elems)
op_sum = tf.foldl(lambda a, x: a + x, elems)
op_mean = tf.reduce_mean(elems)
with tf.Session() as sess:
sess.run(tf.initialize_all_variables())
print sess.run([op_sum, op_mean])
def testFoldShape(self):
with self.test_session():
x = tf.constant([[1, 2, 3], [4, 5, 6]])
def fn(_, current_input):
return current_input
initializer = tf.constant([0, 0, 0])
y = tf.foldl(fn, x, initializer=initializer)
self.assertAllEqual(y.get_shape(), y.eval().shape)
def combined_term_v2(sparse_neighbors,
x_in,
x_out,
features_in,
kernel,
bias=None):
"""
Implements a single term in convolution with feature transform FIRST.
This version is intended for use in UP-sampling, i.e. where
N_in > N_out and/or F_in < F_out.
\\sum_{k,q,j} n_{ij} theta_{kqp} f_{jq} xout_{ki} xin_{kj}
= \\sum_{k} xout_{ki} \\sum_j n_{ij} xin_{kj} \\sum_q theta_{kqp} f_{qj}
For the case where x_out is None (implied to be ones) the sum over k and j
are switched, i.e.
\\\\sum_j n_{ij} sum_{k} xin_{kj} \\sum_q theta_{kqp} f_{qj}.
"""
D, F_in, F_out = _check_arg_shapes(sparse_neighbors, x_in, x_out,
features_in, kernel)
del D, F_in
def get_kth_term(kernel, x_in=None, x_out=None):
fx = features_in
fx = tf.matmul(fx, kernel)
if x_in is not None:
fx = fx * tf.expand_dims(x_in, axis=-1)
if x_out is not None:
# we do sparse multiplication outside foldl if x_out is None
fx = tf.sparse.sparse_dense_matmul(sparse_neighbors, fx)
fx = fx * tf.expand_dims(x_out, axis=-1)
return fx
args = [kernel]
names = ['kernel']
if x_in is not None:
args.append(x_in)
names.append('x_in')
if x_out is not None:
args.append(x_out)
names.append('x_out')
def accumulate(s, args):
return s + get_kth_term(**dict(zip(names, args)))
if bias is None:
init = tf.zeros((sparse_neighbors.dense_shape[0], F_out),
dtype=features_in.dtype)
else:
init = get_kth_term(bias)
if x_out is not None:
init = tf.sparse.sparse_dense_matmul(sparse_neighbors, init)
fx = tf.foldl(accumulate, args, init, parallel_iterations=20)
if x_out is None:
fx = tf.sparse.sparse_dense_matmul(sparse_neighbors, fx)
return fx
def add_all(self, fn, inputs):
initializer = self.accumulate(None, fn(inputs[0]))
inputs = inputs[1:]
@tf.function
def accumulate(acc, input):
return self.accumulate(acc, fn(input))
self.add(tf.foldl(fn=accumulate, elems=inputs,
initializer=initializer))
def dependency_state_indices():
fold = (lambda acc, args: tf.concat(values=(
acc, tf.range(start=args[0], limit=(args[0] + args[1]))),
axis=0))
return tf.foldl(fn=fold,
elems=(subsampled_starts, subsampled_lengths),
initializer=indices[:0],
parallel_iterations=10,
back_prop=False,
swap_memory=False)
def tf_matmul_right(dUs):
"""
Parameters:
dUs: tf.Tensor
Tensorlist of shape (N, n,m)
with number N matrices of size nxm
Multiplies a list of matrices from the right.
"""
return tf.foldl(lambda a, x: tf.matmul(a, x), dUs)
def call(self, x):
vertices, warp_params = x[0], x[1]
vertices_repeated = tf.tile(vertices[tf.newaxis, :, :], [tf.shape(x[1])[0], 1, 1])
warped_vertices = tf.foldl(warp_rbf,
elems=[self.keypoints, self.warp_vectors,
K.permute_dimensions(warp_params, (1, 0, 2))],
initializer=vertices_repeated)
self.warp_strengths = 0.5 - tf.exp(-2.5 * tf.norm(warped_vertices - vertices, axis=2, keepdims=True) ** 2)
return warped_vertices
def build_vocab(word_tensor, vocab_size):
unique, idx = tf.unique(word_tensor)
counts = tf.foldl(lambda counts, item: counts + tf.one_hot(
tf.reshape(item, [-1]), tf.shape(unique)[0], dtype=tf.int32)[0],
idx,
initializer=tf.zeros_like(unique, dtype=tf.int32),
back_prop=False)
_, indices = tf.nn.top_k(counts, k=vocab_size)
return tf.gather(unique, indices)
def _build_loss_psi(self, data):
batch_size = self.data_iterator.shape[0]
batch_zeros = tf.zeros([batch_size])
psi_0 = tf.stack(batch_size * [self.psi_0])
loss = batch_zeros
# We switch to increments
incs = self.data_iterator[:, 1:] - self.data_iterator[:, :-1]
incs = tf.transpose(incs, [1, 0]) # foldl goes along the 1st dimension
_, loss, _ = tf.foldl(self._psi_and_loss_update, incs,
initializer=(psi_0, loss, 0.), name="loss_fold")
return tf.reduce_mean(loss)
def _sample_n(self, n, seed=None):
all_counts = tf.to_float(tf.range(self._total_count + 1))
for batch_dim in range(self.batch_shape.ndims):
all_counts = tf.expand_dims(all_counts, axis=-1)
all_cdfs = tf.map_fn(self.cdf, all_counts)
shape = tf.concat([[n], self.batch_shape_tensor()], 0)
uniform = tf.random_uniform(shape, seed=seed)
return tf.foldl(
lambda acc, cdfs: tf.where(uniform > cdfs, acc + 1, acc),
all_cdfs,
initializer=tf.zeros(shape, dtype=tf.int32))
def testFold_Grad(self):
with self.test_session():
elems = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], name="data")
v = tf.constant(2.0, name="v")
r = tf.foldl(lambda a, x: tf.mul(a, x), elems, initializer=v)
r = tf.gradients(r, v)[0]
self.assertAllEqual(720.0, r.eval())
r = tf.foldr(lambda a, x: tf.mul(a, x), elems, initializer=v)
r = tf.gradients(r, v)[0]
self.assertAllEqual(720.0, r.eval())
def testFold_Grad(self):
with self.test_session():
elems = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], name="data")
v = tf.constant(2.0, name="v")
r = tf.foldl(lambda a, x: tf.mul(a, x), elems, initializer=v)
r = tf.gradients(r, v)[0]
self.assertAllEqual(720.0, r.eval())
r = tf.foldr(lambda a, x: tf.mul(a, x), elems, initializer=v)
r = tf.gradients(r, v)[0]
self.assertAllEqual(720.0, r.eval())
def nested_control_flow_module_fn():
"""Compute the sum of elements greater than 'a' with nested control flow."""
elems = tf.placeholder(
dtype=tf.float32, name="elems", shape=[None])
a = tf.placeholder(dtype=tf.float32, name="a")
def sum_above_a(acc, x):
return acc + tf.cond(x > a, lambda: x, lambda: 0.0)
hub.add_signature(
inputs={"elems": elems, "a": a},
outputs=tf.foldl(sum_above_a, elems, initializer=tf.constant(0.0)))
def batch_sizes_after_windowing(sizes, window):
"""
:param tf.Tensor sizes: (batch_sizes)
:param int window: size of the applied window
:return: sizes for each batch after applying a window on each batch
:rtype: tf.Tensor
"""
def fold_times(acc, x):
r1 = tf.tile([window], [tf.maximum(x - window + 1, 0)])
r2 = tf.range(tf.minimum(x, window - 1), 0, -1)
return tf.concat([acc, r1, r2], 0)
return tf.foldl(fold_times, sizes, tf.placeholder_with_default(tf.zeros([0], dtype='int32'), [None]), name='fold_sizes')
def total_integral(pol):
with tf.name_scope("polyintegral"):
# Non-zero indices
idx = tf.where(tf.not_equal(pol, 0))
# Non-zero values
vals = tf.gather_nd(pol, idx)
ccat = tf.concat(
[tf.expand_dims(vals, 1),
tf.cast(idx, tf.float32)], 1)
return tf.foldl(Polynomial.fold_integral,
ccat,
initializer=tf.constant(0.))
def build_loss_psi(data):
batch_zeros = tf.zeros_like(data[:, 0])
psi_0 = tf.one_hot(tf.cast(batch_zeros, dtype=tf.int32),
self.bond_d,
dtype=tf.complex64)
loss = batch_zeros
data = tf.transpose(data, [1, 0]) # foldl goes along the first dimension
_, loss = tf.foldl(_psi_and_loss_update,
data,
initializer=(psi_0, loss),
name="loss_fold")
return tf.reduce_mean(loss)
def build_vocab(word_tensor, vocab_size):
unique, idx = tf.unique(word_tensor)
counts = tf.foldl(
lambda counts, item: counts + tf.one_hot(
tf.reshape(item, [-1]),
tf.shape(unique)[0],
dtype=tf.int32)[0],
idx,
initializer=tf.zeros_like(unique, dtype=tf.int32),
back_prop=False
)
_, indices = tf.nn.top_k(counts, k=vocab_size)
return tf.gather(unique, indices)
def _log_prob(self, value):
with tf.control_dependencies(self._runtime_assertions):
# The argument `value` is a tensor of sequences of observations.
# `observation_batch_shape` is the shape of that tensor with the
# sequence part removed.
# `observation_batch_shape` is then broadcast to the full batch shape
# to give the `working_shape` that defines the shape of the result.
observation_batch_shape = tf.shape(
value)[:-1 - self._underlying_event_rank]
# value :: observation_batch_shape num_steps observation_event_shape
working_shape = tf.broadcast_dynamic_shape(observation_batch_shape,
self.batch_shape_tensor())
log_init = tf.broadcast_to(self._log_init,
tf.concat([working_shape,
[self._num_states]], axis=0))
# log_init :: working_shape num_states
log_transition = self._log_trans
# `observation_event_shape` is the shape of each sequence of observations
# emitted by the model.
observation_event_shape = tf.shape(
value)[-1 - self._underlying_event_rank:]
working_obs = tf.broadcast_to(value,
tf.concat([working_shape,
observation_event_shape],
axis=0))
# working_obs :: working_shape observation_event_shape
r = self._underlying_event_rank
# Move index into sequence of observations to front so we can apply
# tf.foldl
working_obs = util.move_dimension(working_obs,
-1 - r, 0)[..., tf.newaxis]
# working_obs :: num_steps working_shape underlying_event_shape
observation_probs = (
self._observation_distribution.log_prob(working_obs))
def forward_step(log_prev_step, log_observation):
return _log_vector_matrix(log_prev_step,
log_transition) + log_observation
fwd_prob = tf.foldl(forward_step, observation_probs, initializer=log_init)
# fwd_prob :: working_shape num_states
log_prob = tf.reduce_logsumexp(fwd_prob, axis=-1)
# log_prob :: working_shape
return log_prob
def batch_indices_after_windowing(sizes, window):
"""
here we compute the start and end times for each of the new batches when applying a window
:param tf.Tensor sizes: (batch_sizes)
:param int window: size of the applied window
:return: tensor of shape (?, 3), contains batch index, start-frame and end-frame for each batch after applying a window
:rtype: tf.Tensor
"""
def fold_batches(acc, x):
b = x[0]
l = x[1]
batch = tf.tile([b], [l])
start = tf.range(l)
end = tf.minimum(tf.range(window, l + window), l)
return tf.concat([acc, tf.transpose(tf.stack([batch, start, end]))], axis=0)
return tf.foldl(fold_batches, tf.stack([tf.range(tf.shape(sizes)[0]), sizes], axis=1),
tf.placeholder_with_default(tf.zeros([0, 3], dtype='int32'), [None, 3]), name="fold_batches")
def neural_gpu_body(inputs, hparams, name=None):
"""The core Neural GPU."""
with tf.variable_scope(name, "neural_gpu"):
def step(state, inp): # pylint: disable=missing-docstring
x = tf.nn.dropout(state, 1.0 - hparams.dropout)
for layer in xrange(hparams.num_hidden_layers):
x = common_layers.conv_gru(
x, (hparams.kernel_height, hparams.kernel_width),
hparams.hidden_size,
name="cgru_%d" % layer)
# Padding input is zeroed-out in the modality, we check this by summing.
padding_inp = tf.less(tf.reduce_sum(tf.abs(inp), axis=[1, 2]), 0.00001)
new_state = tf.where(padding_inp, state, x) # No-op where inp is padding.
return new_state
return tf.foldl(
step,
tf.transpose(inputs, [1, 0, 2, 3]),
initializer=inputs,
parallel_iterations=1,
swap_memory=True)
def reduce(self, inputs): return tf.foldl(lambda a, x: a * x, inputs)
def _greedy_infer(self, features, decode_length, last_position_only):
"""A slow greedy inference method.
Quadratic time in decode_length.
Args:
features: an map of string to `Tensor`
decode_length: an integer. How many additional timesteps to decode.
last_position_only: a boolean, speed-up by computing last position only.
Returns:
samples: an integer `Tensor`.
"""
if not features:
features = {}
inputs_old = None
if "inputs" in features and len(features["inputs"].shape) < 4:
inputs_old = features["inputs"]
features["inputs"] = tf.expand_dims(features["inputs"], 2)
if not self.has_input:
features["partial_targets"] = tf.to_int64(features["inputs"])
def infer_step(recent_output, _):
"""Inference step."""
recent_output.set_shape([None, None, None, 1])
padded = tf.pad(recent_output, [[0, 0], [0, 1], [0, 0], [0, 0]])
features["targets"] = padded
# This is inefficient in that it generates samples at all timesteps,
# not just the last one, except if last_position_only is set (dangerous).
samples = self.sample(features, last_position_only=last_position_only)
# Concatenate the already-generated recent_output with last timestep
# of the newly-generated samples.
if last_position_only:
cur_sample = samples[:, -1, :, :]
else:
cur_sample = samples[:, tf.shape(recent_output)[1], :, :]
cur_sample = tf.to_int64(tf.expand_dims(cur_sample, axis=1))
samples = tf.concat([recent_output, cur_sample], axis=1)
samples.set_shape([None, None, None, 1])
return samples
# Create an initial output tensor. This will be passed
# to the infer_step, which adds one timestep at every iteration.
if "partial_targets" in features:
initial_output = tf.convert_to_tensor(features["partial_targets"])
else:
batch_size = tf.shape(features["inputs"])[0]
initial_output = tf.zeros((batch_size, 0, 1, 1), dtype=tf.int64)
# Hack: foldl complains when the output shape is less specified than the
# input shape, so we confuse it about the input shape.
initial_output = tf.slice(initial_output, [0, 0, 0, 0],
tf.shape(initial_output))
if isinstance(self._hparams.problems[self._problem_idx].target_modality,
modality.ClassLabelModality):
decode_length = 1
else:
decode_length = tf.shape(features["inputs"])[1] + decode_length
result = tf.foldl(
infer_step,
tf.range(decode_length),
initializer=initial_output,
back_prop=False,
parallel_iterations=1)
if inputs_old is not None: # Restore to not confuse Estimator.
features["inputs"] = inputs_old
return result
def r_inner(a, x): return tf.foldl(lambda b, y: b * y * x, inner_elems, initializer=a)
