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 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 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)