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)
