# Buggy code
def _call(self, features, labels, params, mode):
    training = mode == tf.estimator.ModeKeys.TRAIN
    self.examples_inputter.build()

    features_length = self.features_inputter.get_length(features)
    source_inputs = self.features_inputter.make_inputs(features, training=training)
    with tf.variable_scope("encoder"):
        encoder_outputs, encoder_state, encoder_sequence_length = self.encoder.encode(
            source_inputs, sequence_length=features_length, mode=mode
        )

    target_vocab_size = self.labels_inputter.vocabulary_size
    target_dtype = self.labels_inputter.dtype
    output_layer = None
    if EmbeddingsSharingLevel.share_target_embeddings(self.share_embeddings):
        output_layer = layers.Dense(
            target_vocab_size,
            weight=self.labels_inputter.embedding,
            transpose=True,
            dtype=target_dtype,
        )
        with tf.name_scope(tf.get_variable_scope().name + "/"):
            output_layer.build([None, self.decoder.output_size])

    if labels is not None:
        target_inputs = self.labels_inputter.make_inputs(labels, training=training)
        with tf.variable_scope("decoder"):
            sampling_probability = None
            if mode == tf.estimator.ModeKeys.TRAIN:
                sampling_probability = get_sampling_probability(
                    tf.train.get_or_create_global_step(),
                    read_probability=params.get("scheduled_sampling_read_probability"),
                    schedule_type=params.get("scheduled_sampling_type"),
                    k=params.get("scheduled_sampling_k"),
                )

            logits, _, _, attention = self.decoder.decode(
                target_inputs,
                self.labels_inputter.get_length(labels),
                vocab_size=target_vocab_size,
                initial_state=encoder_state,
                sampling_probability=sampling_probability,
                embedding=self.labels_inputter.embedding,
                output_layer=output_layer,
                mode=mode,
                memory=encoder_outputs,
                memory_sequence_length=encoder_sequence_length,
                return_alignment_history=True,
            )
            if "alignment" in labels:
                outputs = {"logits": logits, "attention": attention}
            else:
                outputs = logits
    else:
        outputs = None

    if mode != tf.estimator.ModeKeys.TRAIN:
        with tf.variable_scope("decoder", reuse=labels is not None):
            batch_size = tf.shape(
                tf.contrib.framework.nest.flatten(encoder_outputs)[0]
            )[0]
            beam_width = params.get("beam_width", 1)
            maximum_iterations = params.get("maximum_iterations", 250)
            minimum_length = params.get("minimum_decoding_length", 0)
            sample_from = params.get("sampling_topk", 1)
            sample_temperature = params.get("sampling_temperature", 1)
            start_tokens = tf.fill([batch_size], constants.START_OF_SENTENCE_ID)
            end_token = constants.END_OF_SENTENCE_ID

            if beam_width <= 1:
                sampled_ids, _, sampled_length, log_probs, alignment = (
                    self.decoder.dynamic_decode(
                        self.labels_inputter.embedding,
                        start_tokens,
                        end_token,
                        vocab_size=target_vocab_size,
                        initial_state=encoder_state,
                        output_layer=output_layer,
                        maximum_iterations=maximum_iterations,
                        minimum_length=minimum_length,
                        mode=mode,
                        memory=encoder_outputs,
                        memory_sequence_length=encoder_sequence_length,
                        dtype=target_dtype,
                        return_alignment_history=True,
                        sample_from=sample_from,
                        sample_temperature=sample_temperature,
                    )
                )
            else:
                length_penalty = params.get("length_penalty", 0)
                sampled_ids, _, sampled_length, log_probs, alignment = (
                    self.decoder.dynamic_decode_and_search(
                        self.labels_inputter.embedding,
                        start_tokens,
                        end_token,
                        vocab_size=target_vocab_size,
                        initial_state=encoder_state,
                        output_layer=output_layer,
                        beam_width=beam_width,
                        length_penalty=length_penalty,
                        maximum_iterations=maximum_iterations,
                        minimum_length=minimum_length,
                        mode=mode,
                        memory=encoder_outputs,
                        memory_sequence_length=encoder_sequence_length,
                        dtype=target_dtype,
                        return_alignment_history=True,
                        sample_from=sample_from,
                        sample_temperature=sample_temperature,
                    )
                )

        target_vocab_rev = self.labels_inputter.vocabulary_lookup_reverse()
        target_tokens = target_vocab_rev.lookup(tf.cast(sampled_ids, tf.int64))

        if params.get("replace_unknown_target", False):
            if alignment is None:
                raise TypeError(
                    "replace_unknown_target is not compatible with decoders "
                    "that don't return alignment history"
                )
            if not isinstance(self.features_inputter, inputters.WordEmbedder):
                raise TypeError(
                    "replace_unknown_target is only defined when the source "
                    "inputter is a WordEmbedder"
                )
            source_tokens = features["tokens"]
            if beam_width > 1:
                source_tokens = tf.contrib.seq2seq.tile_batch(
                    source_tokens, multiplier=beam_width
                )
            # Merge batch and beam dimensions.
            original_shape = tf.shape(target_tokens)
            target_tokens = tf.reshape(target_tokens, [-1, original_shape[-1]])
            attention = tf.reshape(
                alignment, [-1, tf.shape(alignment)[2], tf.shape(alignment)[3]]
            )
            # We don't have attention for 
 but ensure that the attention time dimension matches
            # the tokens time dimension.
            attention = reducer.align_in_time(attention, tf.shape(target_tokens)[1])
            replaced_target_tokens = replace_unknown_target(
                target_tokens, source_tokens, attention
            )
            target_tokens = tf.reshape(replaced_target_tokens, original_shape)

        predictions = {
            "tokens": target_tokens,
            "length": sampled_length,
            "log_probs": log_probs,
        }
        if alignment is not None:
            predictions["alignment"] = alignment
    else:
        predictions = None

    return outputs, predictions


# Traceback (tensorflow.python.framework.errors_impl.InvalidArgumentError)
INFO:tensorflow:Using parameters:
data:
eval_features_file: dev_mini.txt
train_features_file: train_full.txt
vocabulary: tokens_v3.gpt2
eval:
batch_size: 32
eval_delay: 1200
exporter: last
exporters: last
external_evaluators: null
num_threads: 1
prefetch_buffer_size: 1
save_eval_predictions: false
infer:
batch_size: 32
bucket_width: 5
n_best: 1
num_threads: 1
prefetch_buffer_size: 1
with_scores: false
model_dir: run
params:
average_loss_in_time: true
beam_width: 4
decay_params:
max_step: 1000000
model_dim: 512
warmup_steps: 8000
decay_type: noam_decay_v2
gradients_accum: 1
label_smoothing: 0.1
learning_rate: 2.0
length_penalty: 0.6
maximum_iterations: 50
minimum_decoding_length: 0
optimizer: LazyAdamOptimizer
optimizer_params:
beta1: 0.9
beta2: 0.998
weight_decay: 0.01
score:
batch_size: 64
train:
average_last_checkpoints: 8
batch_size: 4096
batch_type: tokens
bucket_width: 8
effective_batch_size: null
keep_checkpoint_max: 8
maximum_features_length: 256
maximum_labels_length: 256
num_threads: 8
prefetch_buffer_size: 4
sample_buffer_size: 100000
save_checkpoints_secs: null
save_checkpoints_steps: 5000
save_summary_steps: 100
train_steps: 1000000

2019-03-06 23:09:26.843446: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
2019-03-06 23:09:27.382716: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1432] Found device 0 with properties:
name: Tesla K40c major: 3 minor: 5 memoryClockRate(GHz): 0.745
pciBusID: 0000:02:00.0
totalMemory: 11.17GiB freeMemory: 11.09GiB
2019-03-06 23:09:27.645287: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1432] Found device 1 with properties:
name: Tesla K40c major: 3 minor: 5 memoryClockRate(GHz): 0.745
pciBusID: 0000:04:00.0
totalMemory: 11.17GiB freeMemory: 11.09GiB
2019-03-06 23:09:27.910482: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1432] Found device 2 with properties:
name: Tesla K40c major: 3 minor: 5 memoryClockRate(GHz): 0.745
pciBusID: 0000:83:00.0
totalMemory: 11.17GiB freeMemory: 11.09GiB
2019-03-06 23:09:28.184100: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1432] Found device 3 with properties:
name: Tesla K40c major: 3 minor: 5 memoryClockRate(GHz): 0.745
pciBusID: 0000:84:00.0
totalMemory: 11.17GiB freeMemory: 11.09GiB
2019-03-06 23:09:28.184971: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1511] Adding visible gpu devices: 0, 1, 2, 3
2019-03-06 23:09:35.080283: I tensorflow/core/common_runtime/gpu/gpu_device.cc:982] Device interconnect StreamExecutor with strength 1 edge matrix:
2019-03-06 23:09:35.080335: I tensorflow/core/common_runtime/gpu/gpu_device.cc:988]      0 1 2 3
2019-03-06 23:09:35.080348: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1001] 0:   N Y N N
2019-03-06 23:09:35.080356: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1001] 1:   Y N N N
2019-03-06 23:09:35.080364: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1001] 2:   N N N Y
2019-03-06 23:09:35.080372: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1001] 3:   N N Y N
2019-03-06 23:09:35.081500: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/device:GPU:0 with 10747 MB memory) -> physical GPU (device: 0, name: Tesla K40c, pci bus id: 0000:02:00.0, compute capability: 3.5)
2019-03-06 23:09:35.082024: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/device:GPU:1 with 10747 MB memory) -> physical GPU (device: 1, name: Tesla K40c, pci bus id: 0000:04:00.0, compute capability: 3.5)
2019-03-06 23:09:35.082351: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/device:GPU:2 with 10747 MB memory) -> physical GPU (device: 2, name: Tesla K40c, pci bus id: 0000:83:00.0, compute capability: 3.5)
2019-03-06 23:09:35.082622: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1115] Created TensorFlow device (/device:GPU:3 with 10747 MB memory) -> physical GPU (device: 3, name: Tesla K40c, pci bus id: 0000:84:00.0, compute capability: 3.5)
INFO:tensorflow:Using config: {'_model_dir': '/mnt/yardcephfs/mmyard/g_wxg_td_prc/chriscxyan/OpenNMT-TF/opennmt-tf/code/data/gpt2/run', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': 5000, '_save_checkpoints_secs': None, '_session_config': gpu_options {
}
allow_soft_placement: true
graph_options {
rewrite_options {
layout_optimizer: OFF
}
}
, '_keep_checkpoint_max': 8, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_device_fn': None, '_protocol': None, '_eval_distribute': None, '_experimental_distribute': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f079dd36cc0>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}
INFO:tensorflow:Not using Distribute Coordinator.
INFO:tensorflow:Running training and evaluation locally (non-distributed).
INFO:tensorflow:Start train and evaluate loop. The evaluate will happen after every checkpoint. Checkpoint frequency is determined based on RunConfig arguments: save_checkpoints_steps 5000 or save_checkpoints_secs None.
INFO:tensorflow:Calling model_fn.
Traceback (most recent call last):
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 1628, in _create_c_op
c_op = c_api.TF_FinishOperation(op_desc)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Duplicate node name in graph: 'lm/w_embs'

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "/mnt/yardcephfs/mmyard/g_wxg_td_prc/chriscxyan/OpenNMT-TF/opennmt-tf/code/../OpenNMT-tf-master/opennmt/bin/main.py", line 201, in <module>
main()
File "/mnt/yardcephfs/mmyard/g_wxg_td_prc/chriscxyan/OpenNMT-TF/opennmt-tf/code/../OpenNMT-tf-master/opennmt/bin/main.py", line 172, in main
runner.train_and_evaluate(checkpoint_path=args.checkpoint_path)
File "/mnt/yardcephfs/mmyard/g_wxg_td_prc/chriscxyan/OpenNMT-TF/opennmt-tf/OpenNMT-tf-master/opennmt/runner.py", line 295, in train_and_evaluate
result = tf.estimator.train_and_evaluate(estimator, train_spec, eval_spec)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/estimator/training.py", line 471, in train_and_evaluate
return executor.run()
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/estimator/training.py", line 610, in run
return self.run_local()
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/estimator/training.py", line 711, in run_local
saving_listeners=saving_listeners)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 354, in train
loss = self._train_model(input_fn, hooks, saving_listeners)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 1207, in _train_model
return self._train_model_default(input_fn, hooks, saving_listeners)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 1237, in _train_model_default
features, labels, model_fn_lib.ModeKeys.TRAIN, self.config)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 1195, in _call_model_fn
model_fn_results = self._model_fn(features=features, **kwargs)
File "/mnt/yardcephfs/mmyard/g_wxg_td_prc/chriscxyan/OpenNMT-TF/opennmt-tf/OpenNMT-tf-master/opennmt/estimator.py", line 162, in _fn
_loss_op, local_model, features_shards, labels_shards, params, mode)
File "/mnt/yardcephfs/mmyard/g_wxg_td_prc/chriscxyan/OpenNMT-TF/opennmt-tf/OpenNMT-tf-master/opennmt/utils/parallel.py", line 151, in __call__
outputs.append(funs[i](*args[i], **kwargs[i]))
File "/mnt/yardcephfs/mmyard/g_wxg_td_prc/chriscxyan/OpenNMT-TF/opennmt-tf/OpenNMT-tf-master/opennmt/estimator.py", line 231, in _loss_op
logits, _ = model(features, labels, params, mode)
File "/mnt/yardcephfs/mmyard/g_wxg_td_prc/chriscxyan/OpenNMT-TF/opennmt-tf/OpenNMT-tf-master/opennmt/models/model.py", line 85, in __call__
return self._call(features, labels, params, mode)
File "/mnt/yardcephfs/mmyard/g_wxg_td_prc/chriscxyan/OpenNMT-TF/opennmt-tf/OpenNMT-tf-master/opennmt/models/language_model.py", line 56, in _call
self.examples_inputter.build()
File "/mnt/yardcephfs/mmyard/g_wxg_td_prc/chriscxyan/OpenNMT-TF/opennmt-tf/OpenNMT-tf-master/opennmt/inputters/text_inputter.py", line 414, in build
name=compat.name_from_variable_scope("w_embs"))
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/ops/variables.py", line 183, in __call__
return cls._variable_v1_call(*args, **kwargs)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/ops/variables.py", line 146, in _variable_v1_call
aggregation=aggregation)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/ops/variables.py", line 125, in <lambda>
previous_getter = lambda **kwargs: default_variable_creator(None, **kwargs)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/ops/variable_scope.py", line 2444, in default_variable_creator
expected_shape=expected_shape, import_scope=import_scope)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/ops/variables.py", line 187, in __call__
return super(VariableMetaclass, cls).__call__(*args, **kwargs)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/ops/variables.py", line 1329, in __init__
constraint=constraint)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/ops/variables.py", line 1443, in _init_from_args
name=name)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/ops/state_ops.py", line 77, in variable_op_v2
shared_name=shared_name)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/ops/gen_state_ops.py", line 1357, in variable_v2
shared_name=shared_name, name=name)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/framework/op_def_library.py", line 787, in _apply_op_helper
op_def=op_def)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/util/deprecation.py", line 488, in new_func
return func(*args, **kwargs)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 3274, in create_op
op_def=op_def)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 1792, in __init__
control_input_ops)
File "/data1/qspace/chriscxyan/anaconda2/envs/tensorflow/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 1631, in _create_c_op
raise ValueError(str(e))
ValueError: Duplicate node name in graph: 'lm/w_embs'

# Fixed code
def _call(self, features, labels, params, mode):
    training = mode == tf.estimator.ModeKeys.TRAIN

    features_length = self.features_inputter.get_length(features)
    source_inputs = self.features_inputter.make_inputs(features, training=training)
    with tf.variable_scope("encoder"):
        encoder_outputs, encoder_state, encoder_sequence_length = self.encoder.encode(
            source_inputs, sequence_length=features_length, mode=mode
        )

    target_vocab_size = self.labels_inputter.vocabulary_size
    target_dtype = self.labels_inputter.dtype
    if labels is not None:
        target_inputs = self.labels_inputter.make_inputs(labels, training=training)
        with tf.variable_scope("decoder"):
            sampling_probability = None
            if mode == tf.estimator.ModeKeys.TRAIN:
                sampling_probability = get_sampling_probability(
                    tf.train.get_or_create_global_step(),
                    read_probability=params.get("scheduled_sampling_read_probability"),
                    schedule_type=params.get("scheduled_sampling_type"),
                    k=params.get("scheduled_sampling_k"),
                )

            logits, _, _, attention = self.decoder.decode(
                target_inputs,
                self.labels_inputter.get_length(labels),
                vocab_size=target_vocab_size,
                initial_state=encoder_state,
                sampling_probability=sampling_probability,
                embedding=self.labels_inputter.embedding,
                output_layer=self.output_layer,
                mode=mode,
                memory=encoder_outputs,
                memory_sequence_length=encoder_sequence_length,
                return_alignment_history=True,
            )
            if "alignment" in labels:
                outputs = {"logits": logits, "attention": attention}
            else:
                outputs = logits
    else:
        outputs = None

    if mode != tf.estimator.ModeKeys.TRAIN:
        with tf.variable_scope("decoder", reuse=labels is not None):
            batch_size = tf.shape(
                tf.contrib.framework.nest.flatten(encoder_outputs)[0]
            )[0]
            beam_width = params.get("beam_width", 1)
            maximum_iterations = params.get("maximum_iterations", 250)
            minimum_length = params.get("minimum_decoding_length", 0)
            sample_from = params.get("sampling_topk", 1)
            sample_temperature = params.get("sampling_temperature", 1)
            start_tokens = tf.fill([batch_size], constants.START_OF_SENTENCE_ID)
            end_token = constants.END_OF_SENTENCE_ID

            if beam_width <= 1:
                sampled_ids, _, sampled_length, log_probs, alignment = (
                    self.decoder.dynamic_decode(
                        self.labels_inputter.embedding,
                        start_tokens,
                        end_token,
                        vocab_size=target_vocab_size,
                        initial_state=encoder_state,
                        output_layer=self.output_layer,
                        maximum_iterations=maximum_iterations,
                        minimum_length=minimum_length,
                        mode=mode,
                        memory=encoder_outputs,
                        memory_sequence_length=encoder_sequence_length,
                        dtype=target_dtype,
                        return_alignment_history=True,
                        sample_from=sample_from,
                        sample_temperature=sample_temperature,
                    )
                )
            else:
                length_penalty = params.get("length_penalty", 0)
                sampled_ids, _, sampled_length, log_probs, alignment = (
                    self.decoder.dynamic_decode_and_search(
                        self.labels_inputter.embedding,
                        start_tokens,
                        end_token,
                        vocab_size=target_vocab_size,
                        initial_state=encoder_state,
                        output_layer=self.output_layer,
                        beam_width=beam_width,
                        length_penalty=length_penalty,
                        maximum_iterations=maximum_iterations,
                        minimum_length=minimum_length,
                        mode=mode,
                        memory=encoder_outputs,
                        memory_sequence_length=encoder_sequence_length,
                        dtype=target_dtype,
                        return_alignment_history=True,
                        sample_from=sample_from,
                        sample_temperature=sample_temperature,
                    )
                )

        target_vocab_rev = self.labels_inputter.vocabulary_lookup_reverse()
        target_tokens = target_vocab_rev.lookup(tf.cast(sampled_ids, tf.int64))

        if params.get("replace_unknown_target", False):
            if alignment is None:
                raise TypeError(
                    "replace_unknown_target is not compatible with decoders "
                    "that don't return alignment history"
                )
            if not isinstance(self.features_inputter, inputters.WordEmbedder):
                raise TypeError(
                    "replace_unknown_target is only defined when the source "
                    "inputter is a WordEmbedder"
                )
            source_tokens = features["tokens"]
            if beam_width > 1:
                source_tokens = tf.contrib.seq2seq.tile_batch(
                    source_tokens, multiplier=beam_width
                )
            # Merge batch and beam dimensions.
            original_shape = tf.shape(target_tokens)
            target_tokens = tf.reshape(target_tokens, [-1, original_shape[-1]])
            attention = tf.reshape(
                alignment, [-1, tf.shape(alignment)[2], tf.shape(alignment)[3]]
            )
            # We don't have attention for 
 but ensure that the attention time dimension matches
            # the tokens time dimension.
            attention = reducer.align_in_time(attention, tf.shape(target_tokens)[1])
            replaced_target_tokens = replace_unknown_target(
                target_tokens, source_tokens, attention
            )
            target_tokens = tf.reshape(replaced_target_tokens, original_shape)

        predictions = {
            "tokens": target_tokens,
            "length": sampled_length,
            "log_probs": log_probs,
        }
        if alignment is not None:
            predictions["alignment"] = alignment
    else:
        predictions = None

    return outputs, predictions