def inner_loop(i, alive_seq, alive_log_probs, finished_seq,
finished_scores, finished_flags, *states):
"""Inner beam search loop.
There are three groups of tensors, alive, finished, and topk.
The alive group contains information about the current alive sequences
The topk group contains information about alive + topk current decoded words
the finished group contains information about finished sentences, that is,
the ones that have decoded to <EOS>. These are what we return.
The general beam search algorithm is as follows:
While we haven't terminated (pls look at termination condition)
1. Grow the current alive to get beam*2 topk sequences
2. Among the topk, keep the top beam_size ones that haven't reached EOS
into alive
3. Among the topk, keep the top beam_size ones have reached EOS into
finished
Repeat
To make things simple with using fixed size tensors, we will end
up inserting unfinished sequences into finished in the beginning. To stop
that we add -ve INF to the score of the unfinished sequence so that when a
true finished sequence does appear, it will have a higher score than all the
unfinished ones.
Args:
i: loop index
alive_seq: Topk sequences decoded so far [batch_size, beam_size, i+1]
alive_log_probs: probabilities of the beams. [batch_size, beam_size]
finished_seq: Current finished sequences.
[batch_size, beam_size, i+1]
finished_scores: scores for each of these sequences.
[batch_size, beam_size]
finished_flags: finished bools for each of these sequences.
[batch_size, beam_size]
*states: mtf Tensors
Returns:
Tuple of
(Incremented loop index
New alive sequences,
Log probs of the alive sequences,
New finished sequences,
Scores of the new finished sequences,
Flags indicating which sequence in finished as reached EOS,
dict of final decoding states)
"""
states = [
mtf.replace_dimensions(state, batch_and_beam_dim,
[batch_dim, beam_dim]) for state in states
]
# Each inner loop, we carry out three steps:
# 1. Get the current topk items.
# 2. Extract the ones that have finished and haven't finished
# 3. Recompute the contents of finished based on scores.
(top2k_seq, top2k_log_probs, top2k_scores, top2k_finished, new_states,
first_selector) = grow_topk(i, alive_seq, alive_log_probs, states)
with tf.variable_scope("grow_alive"):
alive_seq, alive_log_probs, _, second_selector = grow_alive(
top2k_seq, top2k_scores, top2k_log_probs, top2k_finished)
with tf.variable_scope("grow_finished"):
finished_seq, finished_scores, finished_flags, _ = grow_finished(
finished_seq, finished_scores, finished_flags, top2k_seq,
top2k_scores, top2k_finished)
old_beam_dim = mtf.Dimension("old_beam", beam_dim.size)
selector = mtf.einsum([
mtf.rename_dimension(first_selector, beam_dim.name,
old_beam_dim.name), second_selector
],
output_shape=[batch_dim, old_beam_dim, beam_dim])
gathered_states = []
if use_tpu and layout is not None and mesh_shape is not None:
# This hack combines the beam dimension with some of the batch dimension.
# It makes gathering faster on TPU.
#
# Instead of multiplying by a [beam, beam] selector matrix, we instead
# multiply by a [minor_batch*beam, minor_batch*beam] selector matrix.
# This is theoretically more FLOPs, but it brings the matrix size closer
# to the magic optimal value of 128.
#
# TODO(noam): file a bug with the XLA team to do this automatically
major_batch_size = mtf.tensor_dim_to_mesh_dim_size(
layout, mesh_shape, batch_dim)
major_batch = mtf.Dimension(batch_dim.name, major_batch_size)
minor_batch = mtf.Dimension("minor_batch",
batch_dim.size // major_batch.size)
old_minor_batch = mtf.Dimension("old_minor_batch",
minor_batch.size)
old_combined = mtf.Dimension("old_combined",
minor_batch.size * beam_dim.size)
combined = mtf.Dimension("new_combined", old_combined.size)
same_minor_batch = mtf.to_float(
mtf.equal(mtf.range(mesh, old_minor_batch, tf.float32),
mtf.range(mesh, minor_batch, tf.float32)))
selector = mtf.reshape(
selector, [major_batch, minor_batch, old_beam_dim, beam_dim])
selector = mtf.einsum([selector, same_minor_batch],
output_shape=[
major_batch, old_minor_batch,
old_beam_dim, minor_batch, beam_dim
],
reduced_dims=[])
selector = mtf.reshape(selector,
[major_batch, old_combined, combined])
for state in new_states:
s = mtf.replace_dimensions(state, [batch_dim, beam_dim],
[major_batch, old_combined])
s = mtf.einsum([s, mtf.cast(selector, state.dtype)],
reduced_dims=[old_combined],
output_shape=mtf.replace_dimensions(
state.shape, [batch_dim, beam_dim],
[major_batch, combined]))
gathered_states.append(
mtf.replace_dimensions(s, [major_batch, combined],
batch_and_beam_dim))
else:
for state in new_states:
state = mtf.einsum([
mtf.rename_dimension(state, beam_dim.name,
old_beam_dim.name),
mtf.cast(selector, state.dtype)
],
reduced_dims=[old_beam_dim],
output_shape=state.shape)
state = mtf.replace_dimensions(state, [batch_dim, beam_dim],
batch_and_beam_dim)
gathered_states.append(state)
return (i + 1, alive_seq, alive_log_probs, finished_seq,
finished_scores, finished_flags) + tuple(gathered_states)
def rename_length_to_memory_length(
x, length_name="length", memory_length_name="memory_length"):
return mtf.rename_dimension(x, length_name, memory_length_name)
def _my_concat(a, b):
a = mtf.rename_dimension(a, "beam", "triple_beam")
b = mtf.rename_dimension(b, "double_beam", "triple_beam")
return mtf.concat([a, b], "triple_beam")
def local_2d_self_attention_spatial_blocks(query_antecedent,
kv_channels,
heads,
memory_h_dim=None,
memory_w_dim=None,
mask_right=False,
master_dtype=tf.float32,
slice_dtype=tf.float32,
name=None):
"""Attention to the source position and a neighborhood to the left or right.
The sequence is divided into blocks of length block_size.
Attention for a given query position can only see memory positions
less than or equal to the query position, in the corresponding block
and the previous block.
Args:
query_antecedent: a mtf.Tensor with shape [batch, num_h_blocks,
num_w_blocks, h_dim, w_dim, io_channels] must have the same size as
query_length, but a different name.
kv_channels: a mtf.Dimension (the size of the key and value vectors)
heads: a mtf.Dimension (the number of heads)
memory_h_dim: mtf Dimension, for the memory height block.
memory_w_dim: mtf Dimension, for the memory width block.
mask_right: bool, flag specifying whether we mask out attention to the right
for the decoder.
master_dtype: a tf.dtype
slice_dtype: a tf.dtype
name: an optional string.
Returns:
a Tensor of shape
[batch, num_h_blocks, num_w_blocks, h_dim, w_dim, io_channels]
Raises:
ValueError: if channels or depth don't match.
"""
with tf.variable_scope(
name, default_name="multihead_attention", values=[query_antecedent]):
h_dim, w_dim, io_channels = query_antecedent.shape.dims[-3:]
batch, num_h_blocks, num_w_blocks = query_antecedent.shape.dims[:3]
wq, wk, wv, wo = multihead_attention_vars(
query_antecedent.mesh, heads, io_channels, kv_channels,
master_dtype, slice_dtype, query_antecedent.dtype)
# Rename dimensions for the memory height and width.
memory_antecedent = mtf.rename_dimension(query_antecedent, h_dim.name,
"memory_" + h_dim.name)
memory_antecedent = mtf.rename_dimension(memory_antecedent, w_dim.name,
"memory_" + w_dim.name)
memory_h_dim, memory_w_dim = memory_antecedent.shape.dims[-3:-1]
# Call einsum over the query and memory to get query q, keys k and values v.
q = mtf.einsum([query_antecedent, wq],
mtf.Shape([
batch, heads, num_h_blocks, num_w_blocks, h_dim, w_dim,
kv_channels
]))
k = mtf.einsum([memory_antecedent, wk],
mtf.Shape([batch, heads, num_h_blocks, num_w_blocks,
memory_h_dim, memory_w_dim, kv_channels]))
v = mtf.einsum([memory_antecedent, wv],
mtf.Shape([batch, heads, num_h_blocks, num_w_blocks,
memory_h_dim, memory_w_dim, kv_channels]))
# Halo exchange for memory blocks.
k, v = local_2d_halo_exchange(k, v, num_h_blocks, memory_h_dim,
num_w_blocks, memory_w_dim, mask_right)
# Calculate the causal mask to avoid peeking into the future. We compute
# this once and reuse it for all blocks since the block_size is known.
mask = None
if mask_right:
mask = attention_bias_local_2d_block(query_antecedent.mesh, h_dim, w_dim,
memory_h_dim, memory_w_dim)
output = dot_product_attention(q, k, v, mask=mask)
return mtf.einsum(
[output, wo],
mtf.Shape(
[batch, num_h_blocks, num_w_blocks, h_dim, w_dim, io_channels]))
