def cond_fn(step_num, prev_ids, *unused_states):
"""Should we run another loop iteration."""
overflow = mtf.equal(step_num, num_steps)
has_eos = mtf.reduce_any(mtf.equal(prev_ids, eos_id),
reduced_dim=length_dim)
all_has_eos = mtf.reduce_all(has_eos)
return mtf.logical_not(mtf.logical_or(overflow, all_has_eos))
def masked_local_attention_1d_incremental(x,
prev_k,
prev_v,
step_num,
master_dtype=None,
slice_dtype=None,
params=None,
name=None):
"""Incremental local self-attention (one decode step).
Incremental version of masked_local_attention_1d()
Args:
x: a mtf.Tensor with shape [batch..., io_channels]
prev_k: mtf.Tensor with shape
[batch..., heads, window_length, kv_channels]
prev_v: mtf.Tensor with shape
[batch..., heads, window_length, kv_channels]
step_num: mtf Scalar with dtype tf.int32
master_dtype: a tf.dtype (deprecated)
slice_dtype: a tf.dtype (deprecated)
params: a quadruple of Tensors (see multihead_attention_params())
name: an optional string.
Returns:
y: A mtf.Tensor with shape [batch..., io_channels]
new_k: mtf.Tensor with shape
[batch..., heads, window_length, kv_channels]
new_v: mtf.Tensor with shape
[batch..., heads, window_length, kv_channels]
Raises:
ValueError: if the dimensions do not match.
"""
batch_dims = x.shape.dims[:-1]
io_channels = x.shape.dims[-1]
heads, window_length, kv_channels = prev_k.shape.dims[-3:]
with tf.variable_scope(name, default_name="masked_local_attention_1d"):
if params is None:
wq, wk, wv, wo = multihead_attention_vars(
x.mesh, heads, io_channels, kv_channels,
master_dtype, slice_dtype, x.dtype)
else:
wq, wk, wv, wo = params
q = mtf.einsum([x, wq], mtf.Shape(batch_dims + [heads, kv_channels]))
k = mtf.einsum([x, wk], mtf.Shape(batch_dims + [heads, kv_channels]))
v = mtf.einsum([x, wv], mtf.Shape(batch_dims + [heads, kv_channels]))
current_position = mtf.equal(
mtf.range(x.mesh, window_length, dtype=tf.int32),
mtf.mod(step_num, window_length.size))
k = mtf.where(current_position, k, prev_k, output_shape=prev_k.shape)
v = mtf.where(current_position, v, prev_v, output_shape=prev_v.shape)
o = dot_product_attention(q, k, v, mask=None)
y = mtf.einsum([o, wo], x.shape)
return y, k, v
def attention_mask_ignore_padding(inputs, dtype=tf.float32):
"""Bias for encoder-decoder attention.
Args:
inputs: a mtf.Tensor with shape [..., length_dim]
dtype: a tf.dtype
Returns:
a mtf.Tensor with shape [..., memory_length_dim]
"""
inputs = rename_length_to_memory_length(inputs)
return mtf.cast(mtf.equal(inputs, 0), dtype) * -1e9
def body_fn(step_num, ids, *states):
"""Body function for greedy decoding.
Args:
step_num: a mtf.Tensor
ids: a mtf.Tensor
*states: additional mtf.Tensors
Returns:
new_step_num, new_ids, *new_states
"""
logits, new_states = logits_fn(step_num, ids, states)
vocab_dim = logits.shape.dims[-1]
new_ids = mtf.sample_with_temperature(logits, vocab_dim, temperature)
if forced_ids is not None:
# force the new ids to equal the partial targets where specified
# (positions where partial_targets contain nonzero values)
forced = mtf.gather(forced_ids, step_num, length_dim)
new_ids = forced + new_ids * mtf.to_int32(mtf.equal(forced, 0))
ids += new_ids * mtf.one_hot(step_num, length_dim, dtype=tf.int32)
new_step_num = step_num + 1
return [new_step_num, ids] + new_states
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 grow_topk(i, alive_seq, alive_log_probs, states=None):
r"""Inner beam search loop.
This function takes the current alive sequences, and grows them to topk
sequences where k = 2*beam. We use 2*beam because, we could have beam_size
number of sequences that might hit <EOS> and there will be no alive
sequences to continue. With 2*beam_size, this will not happen. This relies
on the assumption the vocab size is > beam size. If this is true, we'll
have at least beam_size non <EOS> extensions if we extract the next top
2*beam words.
Length penalty is given by = (5+len(decode)/6) ^ -\alpha. Pls refer to
https://arxiv.org/abs/1609.08144.
Args:
i: loop index
alive_seq: Topk sequences decoded so far [batch, beam, length]
alive_log_probs: probabilities of these sequences. [batch, beam]
states: optional list of mtf.Tensor
Returns:
Tuple of
(Topk sequences extended by the next word,
The log probs of these sequences,
The scores with length penalty of these sequences,
Flags indicating which of these sequences have finished decoding,
list of transformed decoding states)
"""
logits, new_states = logits_fn(i, alive_seq, states)
batch_dim, beam_dim, vocab_dim = logits.shape.dims
# Convert logits to normalized log probs
candidate_log_probs = mtf.log_softmax(logits, vocab_dim)
# Multiply the probabilities by the current probabilities of the beam.
# (batch_size, beam_size, vocab_size) + (batch_size, beam_size, 1)
log_probs = candidate_log_probs + alive_log_probs
length_penalty = mtf.pow(((5. + mtf.cast(i + 1, logits.dtype)) / 6.),
alpha)
# scores have shape [batch, beam, vocab]
curr_scores = log_probs / length_penalty
# We find the top 2k sequences to make sure we get k alive sequences.
#
# TODO(noam): This is inefficient. We should separately compute the k
# finished sequences (previously alive sequences + EOS), and the top k new
# alive sequences.
double_beam = mtf.Dimension("double_beam", beam_dim.size * 2)
if use_tpu and layout is not None and mesh_shape is not None:
# Do some partial top-k-ing first locally to avoid communication.
# We reshape the logits from:
# [batch, beam, vocab] to
# [batch, beam, major_vocab, minor_vocab]
# We first reduce (locally) across the minor_vocab dimension. This makes
# the thing we need to broadcast smaller.
# This also enables our shortcut of only picking the top num_prefilter
# sequences per beam per major_vocab in the first pass.
major_vocab_size = mtf.tensor_dim_to_mesh_dim_size(
layout, mesh_shape, vocab_dim)
major_vocab = mtf.Dimension(vocab_dim.name, major_vocab_size)
minor_vocab = mtf.Dimension("minor_vocab",
vocab_dim.size // major_vocab_size)
curr_scores = mtf.reshape(
curr_scores, [batch_dim, beam_dim, major_vocab, minor_vocab])
prefilter = mtf.Dimension("prefilter", num_prefilter
or double_beam.size)
# shape = [batch_dim, beam_dim, major_vocab, prefilter]
top_scores, top_minor_vocab_ids = mtf.top_k(
curr_scores, reduced_dim=minor_vocab, k_dim=prefilter)
combined = mtf.Dimension(
"combined", beam_dim.size * major_vocab.size * prefilter.size)
top_scores = mtf.reshape(top_scores, [batch_dim, combined])
top_minor_vocab_ids = mtf.reshape(top_minor_vocab_ids,
[batch_dim, combined])
# shpae = [batch_dim, double_beam]
# ids are indices representing (beam, major_vocab, prefilter)
top_scores, top_combined_ids = mtf.top_k(top_scores,
reduced_dim=combined,
k_dim=double_beam)
top_minor_vocab_ids = mtf.gather(
top_minor_vocab_ids,
top_combined_ids,
combined,
output_shape=[batch_dim, double_beam])
top_beam_index = top_combined_ids // (major_vocab.size *
prefilter.size)
top_combined_ids -= top_beam_index * (major_vocab.size *
prefilter.size)
top_major_vocab_ids = top_combined_ids // prefilter.size
top_combined_ids -= top_major_vocab_ids * prefilter.size
top_ids = top_major_vocab_ids * minor_vocab.size + top_minor_vocab_ids
else:
beam_and_vocab_dim = mtf.Dimension("beam_and_vocab",
beam_dim.size * vocab_dim.size)
flat_shape = mtf.Shape([batch_dim, beam_and_vocab_dim])
# Flatten out (beam_size, vocab_size) probs into a list of possibilities
flat_curr_scores = mtf.reshape(curr_scores,
flat_shape,
name="flatten_scores")
top_scores, top_ids = mtf.top_k(flat_curr_scores,
reduced_dim=beam_and_vocab_dim,
k_dim=double_beam)
# Work out what beam the top probs are in.
top_beam_index = top_ids // vocab_dim.size
top_ids %= vocab_dim.size # Unflatten the ids
# Recovering the log probs because we will need to send them back
top_log_probs = top_scores * length_penalty
selector = mtf.one_hot(top_beam_index, beam_dim, dtype=tf.float32)
def my_gather(tensor):
return mtf.gather(tensor,
top_beam_index,
beam_dim,
output_shape=mtf.Shape([
double_beam if d == beam_dim else d
for d in tensor.shape.dims
]))
# Gather up the most probable 2*beams both for the ids and finished_in_alive
# bools
top_seq = my_gather(alive_seq)
# Append the most probable alive
top_seq += top_ids * mtf.one_hot(i, length_dim, dtype=tf.int32)
top_finished = mtf.equal(top_ids, eos_id)
return (top_seq, top_log_probs, top_scores, top_finished, new_states,
selector)
