# Buggy code
def _build(self, features, labels, params, mode, config=None):
    features_length = self._get_features_length(features)
    log_dir = config.model_dir if config is not None else None

    source_input_scope = self._get_input_scope(default_name="encoder")
    target_input_scope = self._get_input_scope(default_name="decoder")

    source_inputs = _maybe_reuse_embedding_fn(
        lambda ids: self.source_inputter.transform_data(
            ids, mode=mode, log_dir=log_dir
        ),
        scope=source_input_scope,
    )(features)

    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.target_inputter.vocabulary_size
    target_dtype = self.target_inputter.dtype
    target_embedding_fn = _maybe_reuse_embedding_fn(
        lambda ids: self.target_inputter.transform(ids, mode=mode),
        scope=target_input_scope,
    )

    if labels is not None:
        target_inputs = _maybe_reuse_embedding_fn(
            lambda ids: self.target_inputter.transform_data(
                ids, mode=mode, log_dir=log_dir
            ),
            scope=target_input_scope,
        )(labels)

        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._get_labels_length(labels),
                vocab_size=target_vocab_size,
                initial_state=encoder_state,
                sampling_probability=sampling_probability,
                embedding=target_embedding_fn,
                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(encoder_sequence_length)[0]
            beam_width = params.get("beam_width", 1)
            maximum_iterations = params.get("maximum_iterations", 250)
            minimum_length = params.get("minimum_decoding_length", 0)
            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(
                        target_embedding_fn,
                        start_tokens,
                        end_token,
                        vocab_size=target_vocab_size,
                        initial_state=encoder_state,
                        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,
                    )
                )
            else:
                length_penalty = params.get("length_penalty", 0)
                sampled_ids, _, sampled_length, log_probs, alignment = (
                    self.decoder.dynamic_decode_and_search(
                        target_embedding_fn,
                        start_tokens,
                        end_token,
                        vocab_size=target_vocab_size,
                        initial_state=encoder_state,
                        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,
                    )
                )

        target_vocab_rev = tf.contrib.lookup.index_to_string_table_from_file(
            self.target_inputter.vocabulary_file,
            vocab_size=target_vocab_size - self.target_inputter.num_oov_buckets,
            default_value=constants.UNKNOWN_TOKEN,
        )
        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.source_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]]
            )
            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)
2018-11-28 10:58:38.020515: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:344] Starting optimization for grappler item: tf_graph
INFO:tensorflow:Running local_init_op.
2018-11-28 10:58:38.257302: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:344] Starting optimization for grappler item: tf_graph
INFO:tensorflow:Done running local_init_op.
2018-11-28 10:58:38.349261: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:344] Starting optimization for grappler item: tf_graph
2018-11-28 10:58:38.580481: I tensorflow/core/grappler/optimizers/meta_optimizer.cc:344] Starting optimization for grappler item: tf_graph
Traceback (most recent call last):
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1292, in _do_call
return fn(*args)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1277, in _run_fn
options, feed_dict, fetch_list, target_list, run_metadata)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1365, in _call_tf_sessionrun
run_metadata)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Input to reshape is a tensor with 1818624 values, but the requested shape has 3637248
[[{{node transformer/decoder/while/Reshape_70}} = Reshape[T=DT_FLOAT, Tshape=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](transformer/decoder/while/layer_5/multi_head_1/cond/Merge_1, transformer/decoder/while/Reshape_70/shape)]]
[[{{node transformer/Cast/_1585}} = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_1967_transformer/Cast", tensor_type=DT_INT64, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "opennmt/bin/main.py", line 192, in <module>
main()
File "opennmt/bin/main.py", line 175, in main
log_time=args.log_prediction_time)
File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/runner.py", line 329, in infer
hooks=infer_hooks):
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 568, in predict
preds_evaluated = mon_sess.run(predictions)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 671, in run
run_metadata=run_metadata)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 1148, in run
run_metadata=run_metadata)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 1239, in run
raise six.reraise(*original_exc_info)
File "/home/ari/anaconda3/lib/python3.6/site-packages/six.py", line 693, in reraise
raise value
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 1224, in run
return self._sess.run(*args, **kwargs)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 1296, in run
run_metadata=run_metadata)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/monitored_session.py", line 1076, in run
return self._sess.run(*args, **kwargs)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 887, in run
run_metadata_ptr)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1110, in _run
feed_dict_tensor, options, run_metadata)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1286, in _do_run
run_metadata)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py", line 1306, in _do_call
raise type(e)(node_def, op, message)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Input to reshape is a tensor with 1818624 values, but the requested shape has 3637248
[[node transformer/decoder/while/Reshape_70 (defined at /home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/beam_search.py:88)  = Reshape[T=DT_FLOAT, Tshape=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](transformer/decoder/while/layer_5/multi_head_1/cond/Merge_1, transformer/decoder/while/Reshape_70/shape)]]
[[{{node transformer/Cast/_1585}} = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_1967_transformer/Cast", tensor_type=DT_INT64, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

Caused by op 'transformer/decoder/while/Reshape_70', defined at:
File "opennmt/bin/main.py", line 192, in <module>
main()
File "opennmt/bin/main.py", line 175, in main
log_time=args.log_prediction_time)
File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/runner.py", line 329, in infer
hooks=infer_hooks):
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 550, in predict
features, None, model_fn_lib.ModeKeys.PREDICT, self.config)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/estimator/estimator.py", line 1168, in _call_model_fn
model_fn_results = self._model_fn(features=features, **kwargs)
File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/models/model.py", line 152, in _model_fn
_, predictions = self._build(features, labels, params, mode, config=config)
File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/models/sequence_to_sequence.py", line 277, in _build
return_alignment_history=True))
File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/decoders/decoder.py", line 378, in dynamic_decode_and_search
min_decode_length=minimum_length)
File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/beam_search.py", line 613, in beam_search
back_prop=False)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 3274, in while_loop
return_same_structure)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 2994, in BuildLoop
pred, body, original_loop_vars, loop_vars, shape_invariants)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 2929, in _BuildLoop
body_result = body(*packed_vars_for_body)
File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/beam_search.py", line 539, in inner_loop
i, alive_seq, alive_log_probs, states)
File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/beam_search.py", line 428, in grow_topk
lambda t: _unmerge_beam_dim(t, batch_size, beam_size), flat_states)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/util/nest.py", line 347, in map_structure
structure[0], [func(*x) for x in entries])
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/util/nest.py", line 347, in <listcomp>
structure[0], [func(*x) for x in entries])
File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/beam_search.py", line 428, in <lambda>
lambda t: _unmerge_beam_dim(t, batch_size, beam_size), flat_states)
File "/home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/beam_search.py", line 88, in _unmerge_beam_dim
return tf.reshape(tensor, new_shape)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/gen_array_ops.py", line 6296, in reshape
"Reshape", tensor=tensor, shape=shape, name=name)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/op_def_library.py", line 787, in _apply_op_helper
op_def=op_def)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/util/deprecation.py", line 488, in new_func
return func(*args, **kwargs)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 3274, in create_op
op_def=op_def)
File "/home/ari/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 1770, in __init__
self._traceback = tf_stack.extract_stack()

InvalidArgumentError (see above for traceback): Input to reshape is a tensor with 1818624 values, but the requested shape has 3637248
[[node transformer/decoder/while/Reshape_70 (defined at /home/ari/tf/Onmt-14/OpenNMT-tf/opennmt/utils/beam_search.py:88)  = Reshape[T=DT_FLOAT, Tshape=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](transformer/decoder/while/layer_5/multi_head_1/cond/Merge_1, transformer/decoder/while/Reshape_70/shape)]]
[[{{node transformer/Cast/_1585}} = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_1967_transformer/Cast", tensor_type=DT_INT64, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

# Fixed code
def _build(self, features, labels, params, mode, config=None):
    features_length = self._get_features_length(features)
    log_dir = config.model_dir if config is not None else None

    source_input_scope = self._get_input_scope(default_name="encoder")
    target_input_scope = self._get_input_scope(default_name="decoder")

    source_inputs = _maybe_reuse_embedding_fn(
        lambda ids: self.source_inputter.transform_data(
            ids, mode=mode, log_dir=log_dir
        ),
        scope=source_input_scope,
    )(features)

    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.target_inputter.vocabulary_size
    target_dtype = self.target_inputter.dtype
    target_embedding_fn = _maybe_reuse_embedding_fn(
        lambda ids: self.target_inputter.transform(ids, mode=mode),
        scope=target_input_scope,
    )

    if labels is not None:
        target_inputs = _maybe_reuse_embedding_fn(
            lambda ids: self.target_inputter.transform_data(
                ids, mode=mode, log_dir=log_dir
            ),
            scope=target_input_scope,
        )(labels)

        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._get_labels_length(labels),
                vocab_size=target_vocab_size,
                initial_state=encoder_state,
                sampling_probability=sampling_probability,
                embedding=target_embedding_fn,
                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)
            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(
                        target_embedding_fn,
                        start_tokens,
                        end_token,
                        vocab_size=target_vocab_size,
                        initial_state=encoder_state,
                        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,
                    )
                )
            else:
                length_penalty = params.get("length_penalty", 0)
                sampled_ids, _, sampled_length, log_probs, alignment = (
                    self.decoder.dynamic_decode_and_search(
                        target_embedding_fn,
                        start_tokens,
                        end_token,
                        vocab_size=target_vocab_size,
                        initial_state=encoder_state,
                        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,
                    )
                )

        target_vocab_rev = tf.contrib.lookup.index_to_string_table_from_file(
            self.target_inputter.vocabulary_file,
            vocab_size=target_vocab_size - self.target_inputter.num_oov_buckets,
            default_value=constants.UNKNOWN_TOKEN,
        )
        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.source_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]]
            )
            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