Inference on Universal Transformer fails due to while/Pad in a while loop

[tomweingarten]

Description

Attempting to use t2t-decoder with a universal transformer fails.

Environment information

OS: Ubuntu 19 (Also tested on 18 with same result)

$ pip freeze | grep tensor
mesh-tensorflow==0.0.5
tensor2tensor==1.14.0
tensorboard==1.14.0
tensorflow-datasets==1.2.0
tensorflow-estimator==1.14.0
tensorflow-gan==1.0.0.dev0
tensorflow-gpu==1.14.0
tensorflow-metadata==0.14.0
tensorflow-probability==0.7.0

$ python -V
Python 3.7.3

For bugs: reproduction and error logs

# Steps to reproduce:

t2t-datagen --problem=algorithmic_math_two_variables

t2t-trainer --problem=algorithmic_math_two_variables --model=universal_transformer --hparams_set=universal_transformer_tiny

t2t-decoder --problem=algorithmic_math_two_variables --model=universal_transformer --hparams_set=universal_transformer_tiny

# Error logs:
INFO:tensorflow:Transforming body output with symbol_modality_258_128.top
I0922 17:58:20.209241 140278004143936 t2t_model.py:2248] Transforming body output with symbol_modality_258_128.top
INFO:tensorflow:Cannot use 'while/Pad' as input to 'Identity_5' because 'while/Pad' is in a while loop.

Identity_5 while context: None
while/Pad while context: while/while_context

Traceback for Identity_5:
  File "/home/tom/.local/bin/t2t-decoder", line 17, in <module>
    tf.app.run()
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/platform/app.py", line 40, in run
    _run(main=main, argv=argv, flags_parser=_parse_flags_tolerate_undef)
  File "/home/tom/.local/lib/python3.7/site-packages/absl/app.py", line 299, in run
    _run_main(main, args)
  File "/home/tom/.local/lib/python3.7/site-packages/absl/app.py", line 250, in _run_main
    sys.exit(main(argv))
  File "/home/tom/.local/bin/t2t-decoder", line 12, in main
    t2t_decoder.main(argv)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/bin/t2t_decoder.py", line 205, in main
    decode(estimator, hp, decode_hp)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/bin/t2t_decoder.py", line 106, in decode
    checkpoint_path=FLAGS.checkpoint_path)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/decoding.py", line 224, in decode_from_dataset
    checkpoint_path=checkpoint_path)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/decoding.py", line 313, in decode_once
    for num_predictions, prediction in enumerate(predictions):
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 619, in predict
    features, None, ModeKeys.PREDICT, self.config)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1146, in _call_model_fn
    model_fn_results = self._model_fn(features=features, **kwargs)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1415, in wrapping_model_fn
    use_tpu=use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1472, in estimator_model_fn
    return model.estimator_spec_predict(features, use_tpu=use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1708, in estimator_spec_predict
    features[name] = tf.identity(feature)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/util/dispatch.py", line 180, in wrapper
    return target(*args, **kwargs)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/array_ops.py", line 86, in identity
    ret = gen_array_ops.identity(input, name=name)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/gen_array_ops.py", line 4253, in identity
    "Identity", input=input, name=name)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/op_def_library.py", line 788, in _apply_op_helper
    op_def=op_def)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 3616, in create_op
    op_def=op_def)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 2005, in __init__
    self._traceback = tf_stack.extract_stack()

Traceback for while/Pad:
  File "/home/tom/.local/bin/t2t-decoder", line 17, in <module>
    tf.app.run()
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/platform/app.py", line 40, in run
    _run(main=main, argv=argv, flags_parser=_parse_flags_tolerate_undef)
  File "/home/tom/.local/lib/python3.7/site-packages/absl/app.py", line 299, in run
    _run_main(main, args)
  File "/home/tom/.local/lib/python3.7/site-packages/absl/app.py", line 250, in _run_main
    sys.exit(main(argv))
  File "/home/tom/.local/bin/t2t-decoder", line 12, in main
    t2t_decoder.main(argv)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/bin/t2t_decoder.py", line 205, in main
    decode(estimator, hp, decode_hp)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/bin/t2t_decoder.py", line 106, in decode
    checkpoint_path=FLAGS.checkpoint_path)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/decoding.py", line 224, in decode_from_dataset
    checkpoint_path=checkpoint_path)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/decoding.py", line 313, in decode_once
    for num_predictions, prediction in enumerate(predictions):
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 619, in predict
    features, None, ModeKeys.PREDICT, self.config)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1146, in _call_model_fn
    model_fn_results = self._model_fn(features=features, **kwargs)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1415, in wrapping_model_fn
    use_tpu=use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1472, in estimator_model_fn
    return model.estimator_spec_predict(features, use_tpu=use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1696, in estimator_spec_predict
    use_tpu=use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 817, in infer
    top_beams, alpha, use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/models/research/universal_transformer.py", line 276, in _beam_decode
    top_beams, alpha, use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 940, in _beam_decode_slow
    use_tpu=use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/beam_search.py", line 800, in beam_search
    back_prop=False)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/control_flow_ops.py", line 3501, in while_loop
    return_same_structure)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/control_flow_ops.py", line 3012, in BuildLoop
    pred, body, original_loop_vars, loop_vars, shape_invariants)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/control_flow_ops.py", line 2937, in _BuildLoop
    body_result = body(*packed_vars_for_body)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/beam_search.py", line 717, in inner_loop
    i, alive_seq, alive_log_probs, states)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/beam_search.py", line 597, in grow_topk
    flat_logits = symbols_to_logits_fn(flat_ids)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 873, in symbols_to_logits_fn
    ids = tf.pad(ids[:, 1:], [[0, 0], [0, 1], [0, 0], [0, 0]])
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/array_ops.py", line 2367, in pad
    result = gen_array_ops.pad(tensor, paddings, name=name)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/gen_array_ops.py", line 5993, in pad
    "Pad", input=input, paddings=paddings, name=name)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/op_def_library.py", line 788, in _apply_op_helper
    op_def=op_def)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 3616, in create_op
    op_def=op_def)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 2005, in __init__
    self._traceback = tf_stack.extract_stack()

I0922 17:58:20.311590 140278004143936 control_flow_util.py:349] Cannot use 'while/Pad' as input to 'Identity_5' because 'while/Pad' is in a while loop.

Identity_5 while context: None
while/Pad while context: while/while_context

Traceback for Identity_5:
  File "/home/tom/.local/bin/t2t-decoder", line 17, in <module>
    tf.app.run()
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/platform/app.py", line 40, in run
    _run(main=main, argv=argv, flags_parser=_parse_flags_tolerate_undef)
  File "/home/tom/.local/lib/python3.7/site-packages/absl/app.py", line 299, in run
    _run_main(main, args)
  File "/home/tom/.local/lib/python3.7/site-packages/absl/app.py", line 250, in _run_main
    sys.exit(main(argv))
  File "/home/tom/.local/bin/t2t-decoder", line 12, in main
    t2t_decoder.main(argv)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/bin/t2t_decoder.py", line 205, in main
    decode(estimator, hp, decode_hp)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/bin/t2t_decoder.py", line 106, in decode
    checkpoint_path=FLAGS.checkpoint_path)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/decoding.py", line 224, in decode_from_dataset
    checkpoint_path=checkpoint_path)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/decoding.py", line 313, in decode_once
    for num_predictions, prediction in enumerate(predictions):
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 619, in predict
    features, None, ModeKeys.PREDICT, self.config)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1146, in _call_model_fn
    model_fn_results = self._model_fn(features=features, **kwargs)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1415, in wrapping_model_fn
    use_tpu=use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1472, in estimator_model_fn
    return model.estimator_spec_predict(features, use_tpu=use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1708, in estimator_spec_predict
    features[name] = tf.identity(feature)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/util/dispatch.py", line 180, in wrapper
    return target(*args, **kwargs)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/array_ops.py", line 86, in identity
    ret = gen_array_ops.identity(input, name=name)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/gen_array_ops.py", line 4253, in identity
    "Identity", input=input, name=name)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/op_def_library.py", line 788, in _apply_op_helper
    op_def=op_def)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 3616, in create_op
    op_def=op_def)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 2005, in __init__
    self._traceback = tf_stack.extract_stack()

Traceback for while/Pad:
  File "/home/tom/.local/bin/t2t-decoder", line 17, in <module>
    tf.app.run()
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/platform/app.py", line 40, in run
    _run(main=main, argv=argv, flags_parser=_parse_flags_tolerate_undef)
  File "/home/tom/.local/lib/python3.7/site-packages/absl/app.py", line 299, in run
    _run_main(main, args)
  File "/home/tom/.local/lib/python3.7/site-packages/absl/app.py", line 250, in _run_main
    sys.exit(main(argv))
  File "/home/tom/.local/bin/t2t-decoder", line 12, in main
    t2t_decoder.main(argv)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/bin/t2t_decoder.py", line 205, in main
    decode(estimator, hp, decode_hp)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/bin/t2t_decoder.py", line 106, in decode
    checkpoint_path=FLAGS.checkpoint_path)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/decoding.py", line 224, in decode_from_dataset
    checkpoint_path=checkpoint_path)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/decoding.py", line 313, in decode_once
    for num_predictions, prediction in enumerate(predictions):
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 619, in predict
    features, None, ModeKeys.PREDICT, self.config)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1146, in _call_model_fn
    model_fn_results = self._model_fn(features=features, **kwargs)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1415, in wrapping_model_fn
    use_tpu=use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1472, in estimator_model_fn
    return model.estimator_spec_predict(features, use_tpu=use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1696, in estimator_spec_predict
    use_tpu=use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 817, in infer
    top_beams, alpha, use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/models/research/universal_transformer.py", line 276, in _beam_decode
    top_beams, alpha, use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 940, in _beam_decode_slow
    use_tpu=use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/beam_search.py", line 800, in beam_search
    back_prop=False)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/control_flow_ops.py", line 3501, in while_loop
    return_same_structure)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/control_flow_ops.py", line 3012, in BuildLoop
    pred, body, original_loop_vars, loop_vars, shape_invariants)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/control_flow_ops.py", line 2937, in _BuildLoop
    body_result = body(*packed_vars_for_body)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/beam_search.py", line 717, in inner_loop
    i, alive_seq, alive_log_probs, states)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/beam_search.py", line 597, in grow_topk
    flat_logits = symbols_to_logits_fn(flat_ids)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 873, in symbols_to_logits_fn
    ids = tf.pad(ids[:, 1:], [[0, 0], [0, 1], [0, 0], [0, 0]])
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/array_ops.py", line 2367, in pad
    result = gen_array_ops.pad(tensor, paddings, name=name)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/gen_array_ops.py", line 5993, in pad
    "Pad", input=input, paddings=paddings, name=name)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/op_def_library.py", line 788, in _apply_op_helper
    op_def=op_def)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 3616, in create_op
    op_def=op_def)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 2005, in __init__
    self._traceback = tf_stack.extract_stack()

Traceback (most recent call last):
  File "/home/tom/.local/bin/t2t-decoder", line 17, in <module>
    tf.app.run()
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/platform/app.py", line 40, in run
    _run(main=main, argv=argv, flags_parser=_parse_flags_tolerate_undef)
  File "/home/tom/.local/lib/python3.7/site-packages/absl/app.py", line 299, in run
    _run_main(main, args)
  File "/home/tom/.local/lib/python3.7/site-packages/absl/app.py", line 250, in _run_main
    sys.exit(main(argv))
  File "/home/tom/.local/bin/t2t-decoder", line 12, in main
    t2t_decoder.main(argv)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/bin/t2t_decoder.py", line 205, in main
    decode(estimator, hp, decode_hp)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/bin/t2t_decoder.py", line 106, in decode
    checkpoint_path=FLAGS.checkpoint_path)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/decoding.py", line 224, in decode_from_dataset
    checkpoint_path=checkpoint_path)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/decoding.py", line 313, in decode_once
    for num_predictions, prediction in enumerate(predictions):
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 619, in predict
    features, None, ModeKeys.PREDICT, self.config)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow_estimator/python/estimator/estimator.py", line 1146, in _call_model_fn
    model_fn_results = self._model_fn(features=features, **kwargs)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1415, in wrapping_model_fn
    use_tpu=use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1472, in estimator_model_fn
    return model.estimator_spec_predict(features, use_tpu=use_tpu)
  File "/home/tom/.local/lib/python3.7/site-packages/tensor2tensor/utils/t2t_model.py", line 1708, in estimator_spec_predict
    features[name] = tf.identity(feature)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/util/dispatch.py", line 180, in wrapper
    return target(*args, **kwargs)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/array_ops.py", line 86, in identity
    ret = gen_array_ops.identity(input, name=name)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/gen_array_ops.py", line 4253, in identity
    "Identity", input=input, name=name)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/op_def_library.py", line 788, in _apply_op_helper
    op_def=op_def)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/util/deprecation.py", line 507, in new_func
    return func(*args, **kwargs)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 3616, in create_op
    op_def=op_def)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 2043, in __init__
    self._control_flow_post_processing()
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py", line 2052, in _control_flow_post_processing
    control_flow_util.CheckInputFromValidContext(self, input_tensor.op)
  File "/home/tom/.local/lib/python3.7/site-packages/tensorflow/python/ops/control_flow_util.py", line 350, in CheckInputFromValidContext
    raise ValueError(error_msg + " See info log for more details.")
ValueError: Cannot use 'while/Pad' as input to 'Identity_5' because 'while/Pad' is in a while loop. See info log for more details.

[senarvi]

I saw the same error with transformer_moe model with tensor2tensor 1.14.0, but now I don't see it anymore with the head version. Maybe it has been fixed recently?

[senarvi]

I did some bisecting to find the fixing comment: https://github.com/tensorflow/tensor2tensor/commit/448ce0f56d7f67f7b610c1db72e2ea9529c44511