def testEmpty(self): for dtype in [ dtypes.float32, dtypes.float64, dtypes.int32, dtypes.int64, dtypes.bool ]: with self.test_session(use_gpu=True): test_shapes = [(), (1, ), (2, 3), (0, 2), (2, 3, 5), (2, 0, 5)] for shape in test_shapes: val = inplace_ops.empty(shape, dtype).eval() self.assertEqual(val.shape, shape) self.assertEqual(val.dtype, dtype.as_numpy_dtype) val = inplace_ops.empty(shape, dtype, init=True).eval() self.assertEqual(val.shape, shape) self.assertEqual(val.dtype, dtype.as_numpy_dtype) self.assertAllEqual(val, np.zeros(shape, dtype.as_numpy_dtype)) val = inplace_ops.empty_like(array_ops.zeros( shape, dtype)).eval() self.assertEqual(val.shape, shape) self.assertEqual(val.dtype, dtype.as_numpy_dtype) val = inplace_ops.empty_like(array_ops.zeros(shape, dtype), init=True).eval() self.assertEqual(val.shape, shape) self.assertEqual(val.dtype, dtype.as_numpy_dtype) self.assertAllEqual(val, np.zeros(shape, dtype.as_numpy_dtype)) val = inplace_ops.empty((1, 2), dtypes.string, init=True).eval() self.assertEqual(val.tolist(), [[b"", b""]]) val = inplace_ops.empty((1, 2), dtypes.string, init=False).eval() self.assertEqual(val.tolist(), [[b"", b""]])
def testEmpty(self): for dtype in [ dtypes.float32, dtypes.float64, dtypes.int32, dtypes.int64, dtypes.bool, dtypes.uint8 ]: with self.test_session(use_gpu=True): test_shapes = [(), (1,), (2, 3), (0, 2), (2, 3, 5), (2, 0, 5)] for shape in test_shapes: val = inplace_ops.empty(shape, dtype).eval() self.assertEqual(val.shape, shape) self.assertEqual(val.dtype, dtype.as_numpy_dtype) val = inplace_ops.empty(shape, dtype, init=True).eval() self.assertEqual(val.shape, shape) self.assertEqual(val.dtype, dtype.as_numpy_dtype) self.assertAllEqual(val, np.zeros(shape, dtype.as_numpy_dtype)) val = inplace_ops.empty_like(array_ops.zeros(shape, dtype)).eval() self.assertEqual(val.shape, shape) self.assertEqual(val.dtype, dtype.as_numpy_dtype) val = inplace_ops.empty_like( array_ops.zeros(shape, dtype), init=True).eval() self.assertEqual(val.shape, shape) self.assertEqual(val.dtype, dtype.as_numpy_dtype) self.assertAllEqual(val, np.zeros(shape, dtype.as_numpy_dtype)) with self.test_session(use_gpu=True): val = inplace_ops.empty((1, 2), dtypes.string, init=True).eval() self.assertEqual(val.tolist(), [[b"", b""]]) val = inplace_ops.empty((1, 2), dtypes.string, init=False).eval() self.assertEqual(val.tolist(), [[b"", b""]])
def _EmptyWithFixShape(shape, nmap): """Creates a set of empty initialized tensors with fixed shape. Args: shape: A list of integers to describe the output tensor shape. nmap: A `.NestedMap` of tensors. Returns: A `.NestedMap` with the same keys as nmap. ret.key, a tensor, has the same dtype as nmap.key, but with the fixed shape. """ return nmap.Transform(lambda x: inplace_ops.empty(shape, dtype=x.dtype))
def _EmptyAccForTensor(tensor): return inplace_ops.empty( array_ops.concat([[slen], array_ops.shape(tensor)], axis=0), tensor.dtype, init=True)
def Fill(x): return inplace_ops.empty(tf.concat([[slen], tf.shape(x)], axis=0), x.dtype, init=True)
def _scan(fn, elems, initial, reverse=False, inclusive=False, final_only=False): """Repeatedly applies callable `fn` to a sequence of elements. Implemented by functional_ops.While, tpu friendly, no gradient. This is similar to functional_ops.scan but significantly faster on tpu/gpu for the forward backward use case. Examples: scan(lambda a, e: a + e, [1.0, 2.0, 3.0], 1.0) => [2.0, 4.0, 7.0] Multiple accumulators: scan(lambda a, e: (a[0] + e, a[1] * e), [1.0, 2.0, 3.0], (0.0, 1.0)) Multiple inputs: scan(lambda a, e: a + (e[0] * e[1]), (elems1, elems2), 0.0) Args: fn: callable, fn(accumulators, element) return new accumulator values. The (possibly nested) sequence of accumulators is the same as `initial` and the return value must have the same structure. elems: A (possibly nested) tensor which will be unpacked along the first dimension. The resulting slices will be the second argument to fn. The first dimension of all nested input tensors must be the same. initial: A tensor or (possibly nested) sequence of tensors with initial values for the accumulators. reverse: (optional) True enables scan and output elems in reverse order. inclusive: (optional) True includes the initial accumulator values in the output. Length of output will be len(elem sequence) + 1. Not meaningful if final_only is True. final_only: (optional) When True, return only the final accumulated values, not the concatenation of accumulated values for each input. Returns: A (possibly nested) sequence of tensors with the results of applying fn to tensors unpacked from elems and previous accumulator values. """ flat_elems = [ops.convert_to_tensor(x) for x in nest.flatten(elems)] num_elems = array_ops.shape(flat_elems[0])[0] pack_elems = lambda x: nest.pack_sequence_as(structure=elems, flat_sequence=x) flat_initial = [ops.convert_to_tensor(x) for x in nest.flatten(initial)] pack = lambda x: nest.pack_sequence_as(structure=initial, flat_sequence=x) accum_dtypes = [x.dtype for x in flat_initial] num_accums = len(flat_initial) # Types for counter, [outputs], [accumulators] loop arguments. if final_only: loop_dtypes = [dtypes.int32, dtypes.int32] + accum_dtypes else: loop_dtypes = [dtypes.int32, dtypes.int32 ] + accum_dtypes + accum_dtypes # TODO(tombagby): Update to tfe.defun def cond(i, num_elems, *args): del args return i >= 0 if reverse else i < num_elems # The loop *args are [output tensors] + [accumulator tensors] which must # be paired. Each output corresponds to one accumulator. def body(i, num_elems, *args): """Loop body.""" i.set_shape([]) if final_only: accum = args else: out, accum = args[:num_accums], args[num_accums:] slices = [array_ops.gather(e, i) for e in flat_elems] accum = fn(pack(accum), pack_elems(slices)) flat_accum = nest.flatten(accum) if final_only: new_out = [] else: update_i = i + 1 if inclusive and not reverse else i new_out = [ inplace_ops.alias_inplace_update(x, update_i, y) for x, y in zip(out, flat_accum) ] i = i - 1 if reverse else i + 1 return [i, num_elems] + new_out + flat_accum init_i = (array_ops.shape(flat_elems[0])[0] - 1 if reverse else constant_op.constant(0, dtype=dtypes.int32)) outputs = [] if not final_only: num_outputs = array_ops.shape( flat_elems[0])[0] + (1 if inclusive else 0) for initial_accum in flat_initial: out_shape = array_ops.concat( [[num_outputs], array_ops.shape(initial_accum)], 0) out = inplace_ops.empty(out_shape, dtype=initial_accum.dtype, init=True) if inclusive: out = inplace_ops.alias_inplace_add( out, init_i + (1 if reverse else 0), initial_accum) outputs.append(out) loop_in = [init_i, num_elems] + outputs + flat_initial hostmem = [ i for i, x in enumerate(loop_in) if x.dtype.base_dtype in (dtypes.int32, dtypes.int64) ] if context.executing_eagerly(): loop_results = loop_in while cond(*loop_results): loop_results = body(*loop_results) else: # TODO(tombagby): Update to while_v2. cond = function.Defun(*loop_dtypes)(cond) body = function.Defun(*loop_dtypes)(body) loop_results = functional_ops.While(loop_in, cond, body, hostmem=hostmem) out = loop_results[2:num_accums + 2] return pack(out)
def GreedySearchDecode(self, theta, encoder_outputs, init_beam_search_state=None, pre_beam_search_step_callback=None, post_beam_search_step_callback=None, max_steps=None): """Performs greedy-search based decoding. Args: theta: A NestedMap object containing weights' values of the decoder layer and its children layers. encoder_outputs: A NestedMap containing encoder outputs to be passed to the callbacks. init_beam_search_state: The `InitBeamSearchState` callback. Please refer to the class header comments for more details. pre_beam_search_step_callback: The `PreBeamSearchStepCallback` callback. Please refer to the class header comments for more details. post_beam_search_step_callback: The `PostBeamSearchStepCallback` callback. Please refer to the class header comments for more details. max_steps: maximum beam search steps. If None, use self.params.target_seq_len. Returns: A tuple (hyp_ids, hyp_lens, done_hyps). Note that num_hyps is same as src_batch_size. - hyp_ids: [num_hyps, max_step]. Hyps end with <eos> token if the <eos> token is encountered during search. - hyp_lens: [num_hyps]. - done_hyps: [num_hyps], whether or not an eos is encountered. """ p = self.params if max_steps is None: max_steps = p.target_seq_len initial_results, other_states = init_beam_search_state( theta, encoder_outputs, 1 # num_hyps_per_beam ) num_hyps = tf.shape(initial_results.log_probs)[0] if 'step_ids' in initial_results: # [num_hyps, 1] step_ids = tf.ensure_shape(initial_results.step_ids, [None, 1]) else: step_ids = tf.fill([num_hyps, 1], tf.constant(p.target_sos_id, dtype=tf.int32)) cur_step = tf.constant(0, dtype=tf.int32) done_hyps = inplace_ops.empty(shape=[num_hyps], dtype=tf.bool, init=True, name='done_hyps') hyp_lens = inplace_ops.empty(shape=[num_hyps], dtype=tf.int32, init=True, name='hyp_lens') hyp_ids = inplace_ops.empty(shape=[max_steps, num_hyps], dtype=tf.int32, init=True, name='hyp_ids') def LoopContinue(cur_step, unused_step_ids, unused_hyp_ids, unused_hyp_lens, done_hyps, unused_other_states_list): return tf.logical_and(cur_step < max_steps, tf.logical_not(tf.reduce_all(done_hyps))) def LoopBody(cur_step, step_ids, hyp_ids, hyp_lens, done_hyps, other_states_list): (cur_step, new_step_ids, hyp_ids, hyp_lens, done_hyps, new_other_states) = self._GreedySearchStep( theta, encoder_outputs, cur_step, step_ids, hyp_ids, hyp_lens, done_hyps, other_states.Pack(other_states_list), pre_beam_search_step_callback, post_beam_search_step_callback) return (cur_step, new_step_ids, hyp_ids, hyp_lens, done_hyps, new_other_states.Flatten()) flat_other_states = other_states.Flatten() _, _, final_hyp_ids, final_hyp_lens, final_done_hyps, _ = tf.while_loop( LoopContinue, LoopBody, loop_vars=(cur_step, step_ids, hyp_ids, hyp_lens, done_hyps, flat_other_states), parallel_iterations=10, back_prop=False, swap_memory=False, shape_invariants=(tf.TensorShape(cur_step.get_shape()), tf.TensorShape(step_ids.get_shape()), tf.TensorShape(hyp_ids.get_shape()), tf.TensorShape(hyp_lens.get_shape()), tf.TensorShape(done_hyps.get_shape()), _GetShapes(flat_other_states, none_shapes=True))) # transpose hyp_ids so it matches BeamSearchDecode's output final_hyp_ids = tf.transpose(final_hyp_ids) return final_hyp_ids, final_hyp_lens, final_done_hyps
def _scan(fn, elems, initial, reverse=False, inclusive=False, final_only=False): """Repeatedly applies callable `fn` to a sequence of elements. Implemented by functional_ops.While, tpu friendly, no gradient. This is similar to functional_ops.scan but significantly faster on tpu/gpu for the forward backward use case. Examples: scan(lambda a, e: a + e, [1.0, 2.0, 3.0], 1.0) => [2.0, 4.0, 7.0] Multiple accumulators: scan(lambda a, e: (a[0] + e, a[1] * e), [1.0, 2.0, 3.0], (0.0, 1.0)) Multiple inputs: scan(lambda a, e: a + (e[0] * e[1]), (elems1, elems2), 0.0) Args: fn: callable, fn(accumulators, element) return new accumulator values. The (possibly nested) sequence of accumulators is the same as `initial` and the return value must have the same structure. elems: A (possibly nested) tensor which will be unpacked along the first dimension. The resulting slices will be the second argument to fn. The first dimension of all nested input tensors must be the same. initial: A tensor or (possibly nested) sequence of tensors with initial values for the accumulators. reverse: (optional) True enables scan and output elems in reverse order. inclusive: (optional) True includes the initial accumulator values in the output. Length of output will be len(elem sequence) + 1. Not meaningful if final_only is True. final_only: (optional) When True, return only the final accumulated values, not the concatenation of accumulated values for each input. Returns: A (possibly nested) sequence of tensors with the results of applying fn to tensors unpacked from elems and previous accumulator values. """ flat_elems = [ops.convert_to_tensor(x) for x in nest.flatten(elems)] num_elems = array_ops.shape(flat_elems[0])[0] pack_elems = lambda x: nest.pack_sequence_as(structure=elems, flat_sequence=x) flat_initial = [ops.convert_to_tensor(x) for x in nest.flatten(initial)] pack = lambda x: nest.pack_sequence_as(structure=initial, flat_sequence=x) accum_dtypes = [x.dtype for x in flat_initial] num_accums = len(flat_initial) # Types for counter, [outputs], [accumulators] loop arguments. if final_only: loop_dtypes = [dtypes.int32, dtypes.int32] + accum_dtypes else: loop_dtypes = [dtypes.int32, dtypes.int32] + accum_dtypes + accum_dtypes # TODO(tombagby): Update to tfe.defun @function.Defun(*loop_dtypes) def cond(i, num_elems, *args): del args return i >= 0 if reverse else i < num_elems # The loop *args are [output tensors] + [accumulator tensors] which must # be paired. Each output corresponds to one accumulator. @function.Defun(*loop_dtypes) def body(i, num_elems, *args): """Loop body.""" i.set_shape([]) if final_only: accum = args else: out, accum = args[:num_accums], args[num_accums:] slices = [array_ops.gather(e, i) for e in flat_elems] accum = fn(pack(accum), pack_elems(slices)) flat_accum = nest.flatten(accum) if final_only: new_out = [] else: update_i = i + 1 if inclusive and not reverse else i new_out = [inplace_ops.alias_inplace_update(x, update_i, y) for x, y in zip(out, flat_accum)] i = i - 1 if reverse else i + 1 return [i, num_elems] + new_out + flat_accum init_i = (array_ops.shape(flat_elems[0])[0] - 1 if reverse else constant_op.constant(0, dtype=dtypes.int32)) outputs = [] if not final_only: num_outputs = array_ops.shape(flat_elems[0])[0] + (1 if inclusive else 0) for initial_accum in flat_initial: out_shape = array_ops.concat( [[num_outputs], array_ops.shape(initial_accum)], 0) out = inplace_ops.empty(out_shape, dtype=initial_accum.dtype, init=True) if inclusive: out = inplace_ops.alias_inplace_add( out, init_i + (1 if reverse else 0), initial_accum) outputs.append(out) loop_in = [init_i, num_elems] + outputs + flat_initial hostmem = [ i for i, x in enumerate(loop_in) if x.dtype.base_dtype in (dtypes.int32, dtypes.int64) ] # TODO(tombagby): Update to while_v2. loop_results = functional_ops.While(loop_in, cond, body, hostmem=hostmem) out = loop_results[2:num_accums + 2] return pack(out)