Repository: OpenNMT/OpenNMT-tf

Issue:
Mixed Precision error training with RNN models
I am getting an error when running Sequence to Sequence models with mixed precision.
Models I have encountered this issue with:

- NMTSmallV1
- NMTBigV1

The models will work fine if I remove the option  _--mixed_precision_ 
I can also run a Transformer model fine with mixed precision so it seems to only affect the Sequence to Sequence models.

TensorFlow version: 2.4.1
Python version: 3.6.13
CUDA 11.1
I have a git cloned version of the latest OpenNMT-tf version which adds TF 2.5 support

Example command that causes this:
`python opennmt/bin/main.py --model_type NMTBigV1 --config ~/data_bnmt_en_fr.yml --auto_config --mixed_precision train --with_eval --num_gpus 6`

```

This is the error I am getting:

2021-05-25 12:50:32.018136: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-05-25 12:50:32.020761: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-05-25 12:50:32.023276: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-05-25 12:50:32.025797: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-05-25 12:50:32.028300: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-05-25 12:50:32.030800: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-05-25 12:50:32.033302: I tensorflow/compiler/jit/xla_cpu_device.cc:41] Not creating XLA devices, tf_xla_enable_xla_devices not set
2021-05-25 12:50:32.036000: I device_compatibility_check.py:129] Mixed precision compatibility check (mixed_float16): OK
Your GPUs will likely run quickly with dtype policy mixed_float16 as they all have compute capability of at least 7.0
2021-05-25 12:50:33.115000: W runner.py:242] No checkpoint to restore in ../Models/en-fr/CCAligned/fr-en-ccaligned-bignmt-v2
2021-05-25 12:50:33.119000: W deprecation.py:339] From /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/summary/summary_iterator.py:31: tf_record_iterator (from tensorflow.python.lib.io.tf_record) is deprecated and will be removed in a future version.
Instructions for updating:
Use eager execution and: 
`tf.data.TFRecordDataset(path)`
2021-05-25 12:50:36.613000: I mirrored_strategy.py:350] Using MirroredStrategy with devices ('/job:localhost/replica:0/task:0/device:GPU:0', '/job:localhost/replica:0/task:0/device:GPU:1', '/job:localhost/replica:0/task:0/device:GPU:2', '/job:localhost/replica:0/task:0/device:GPU:3', '/job:localhost/replica:0/task:0/device:GPU:4', '/job:localhost/replica:0/task:0/device:GPU:5')
2021-05-25 12:50:42.719000: I dataset_ops.py:1996] Training on 13992542 examples
2021-05-25 12:50:44.009313: I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:116] None of the MLIR optimization passes are enabled (registered 2)
2021-05-25 12:50:44.013304: I tensorflow/core/platform/profile_utils/cpu_utils.cc:112] CPU Frequency: 3616000000 Hz
2021-05-25 12:51:01.734000: E tensor_array_ops.py:1315] Error: Input value Tensor("nmt_big_v1_1/attentional_rnn_decoder_1/while/attention_wrapper_1/LuongAttention/Softmax:0", shape=(None, None), dtype=float16, device=/job:localhost/replica:0/task:0/device:GPU:0) has dtype <dtype: 'float16'>, but expected dtype <dtype: 'float32'>.  This leads to undefined behavior and will be an error in future versions of TensorFlow.  Traceback:
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/threading.py", line 884, in _bootstrap
    self._bootstrap_inner()
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/threading.py", line 916, in _bootstrap_inner
    self.run()
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/distribute/mirrored_run.py", line 323, in run
    self.main_result = self.main_fn(*self.main_args, **self.main_kwargs)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 667, in wrapper
    return converted_call(f, args, kwargs, options=options)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 459, in converted_call
    result = converted_f(*effective_args, **kwargs)
  File "/tmp/tmpx3typ5u2.py", line 11, in tf___forward
    (loss, gradients) = ag__.converted_call(ag__.ld(self)._compute_gradients, (ag__.ld(source), ag__.ld(target), ag__.ld(accum_steps), ag__.ld(report_steps)), None, fscope)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 461, in converted_call
    result = converted_f(*effective_args)
  File "/tmp/tmpu3du9bqd.py", line 67, in tf___compute_gradients
    ag__.if_stmt(ag__.converted_call(ag__.ld(tf).executing_eagerly, (), None, fscope), if_body_2, else_body_2, get_state_2, set_state_2, ('gradients', 'reported_loss'), 2)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/operators/control_flow.py", line 1165, in if_stmt
    _py_if_stmt(cond, body, orelse)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/operators/control_flow.py", line 1218, in _py_if_stmt
    return body() if cond else orelse()
  File "/tmp/tmpu3du9bqd.py", line 61, in else_body_2
    (training_loss, reported_loss) = ag__.converted_call(ag__.ld(self)._run_model, (ag__.ld(source), ag__.ld(target)), dict(accum_steps=ag__.ld(accum_steps)), fscope)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 459, in converted_call
    result = converted_f(*effective_args, **kwargs)
  File "/tmp/tmpx78hsj6g.py", line 12, in tf___run_model
    (outputs, _) = ag__.converted_call(ag__.ld(self)._model, (ag__.ld(source),), dict(labels=ag__.ld(target), training=True, step=ag__.ld(self)._optimizer.iterations), fscope)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 396, in converted_call
    return _call_unconverted(f, args, kwargs, options)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 478, in _call_unconverted
    return f(*args, **kwargs)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py", line 1012, in __call__
    outputs = call_fn(inputs, *args, **kwargs)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 667, in wrapper
    return converted_call(f, args, kwargs, options=options)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 459, in converted_call
    result = converted_f(*effective_args, **kwargs)
  File "/tmp/tmpey7rqdfk.py", line 30, in tf__call
    ag__.if_stmt((ag__.ld(labels) is not None), if_body, else_body, get_state, set_state, ('outputs',), 1)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/operators/control_flow.py", line 1165, in if_stmt
    _py_if_stmt(cond, body, orelse)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/operators/control_flow.py", line 1218, in _py_if_stmt
    return body() if cond else orelse()
  File "/tmp/tmpey7rqdfk.py", line 25, in if_body
    outputs = ag__.converted_call(ag__.ld(self)._decode_target, (ag__.ld(labels), ag__.ld(encoder_outputs), ag__.ld(encoder_state), ag__.ld(encoder_sequence_length)), dict(step=ag__.ld(step), training=ag__.ld(training)), fscope)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 459, in converted_call
    result = converted_f(*effective_args, **kwargs)
  File "/tmp/tmpy6qhbxsm.py", line 31, in tf___decode_target
    (logits, _, attention) = ag__.converted_call(ag__.ld(self).decoder, (ag__.ld(target_inputs), ag__.converted_call(ag__.ld(self).labels_inputter.get_length, (ag__.ld(labels),), None, fscope)), dict(state=ag__.ld(initial_state), input_fn=ag__.ld(input_fn), sampling_probability=ag__.ld(sampling_probability), training=ag__.ld(training)), fscope)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 396, in converted_call
    return _call_unconverted(f, args, kwargs, options)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 478, in _call_unconverted
    return f(*args, **kwargs)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py", line 1012, in __call__
    outputs = call_fn(inputs, *args, **kwargs)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 667, in wrapper
    return converted_call(f, args, kwargs, options=options)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 459, in converted_call
    result = converted_f(*effective_args, **kwargs)
  File "/tmp/tmpgwed_v1y.py", line 75, in tf__call
    ag__.if_stmt((ag__.ld(rank) == 2), if_body_3, else_body_3, get_state_3, set_state_3, ('attention', 'state', 'logits'), 3)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/operators/control_flow.py", line 1165, in if_stmt
    _py_if_stmt(cond, body, orelse)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/operators/control_flow.py", line 1218, in _py_if_stmt
    return body() if cond else orelse()
  File "/tmp/tmpgwed_v1y.py", line 71, in else_body_3
    ag__.if_stmt((ag__.ld(rank) == 3), if_body_2, else_body_2, get_state_2, set_state_2, ('attention', 'state', 'logits'), 3)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/operators/control_flow.py", line 1165, in if_stmt
    _py_if_stmt(cond, body, orelse)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/operators/control_flow.py", line 1218, in _py_if_stmt
    return body() if cond else orelse()
  File "/tmp/tmpgwed_v1y.py", line 64, in if_body_2
    (logits, state, attention) = ag__.converted_call(ag__.ld(self).forward, (ag__.ld(inputs),), dict(sequence_length=ag__.ld(length_or_step), initial_state=ag__.ld(state), memory=ag__.ld(self).memory, memory_sequence_length=ag__.ld(self).memory_sequence_length, input_fn=ag__.ld(input_fn), sampling_probability=ag__.ld(sampling_probability), training=ag__.ld(training)), fscope)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 459, in converted_call
    result = converted_f(*effective_args, **kwargs)
  File "/tmp/tmpvmpbeshd.py", line 146, in tf__forward
    (_, state, _, outputs_ta, attention_ta) = ag__.converted_call(ag__.ld(tf).while_loop, (ag__.autograph_artifact((lambda *arg: True)), ag__.ld(_body)), dict(loop_vars=(ag__.ld(step), ag__.ld(initial_state), ag__.converted_call(ag__.ld(inputs_ta).read, (0,), None, fscope), ag__.ld(outputs_ta), ag__.ld(attention_ta)), parallel_iterations=32, swap_memory=True, maximum_iterations=ag__.ld(max_step)), fscope)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 396, in converted_call
    return _call_unconverted(f, args, kwargs, options)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 478, in _call_unconverted
    return f(*args, **kwargs)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/util/deprecation.py", line 605, in new_func
    return func(*args, **kwargs)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 2499, in while_loop_v2
    return_same_structure=True)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/control_flow_ops.py", line 2696, in while_loop
    back_prop=back_prop)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py", line 200, in while_loop
    add_control_dependencies=add_control_dependencies)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py", line 990, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py", line 178, in wrapped_body
    outputs = body(*_pack_sequence_as(orig_loop_vars, args))
  File "/tmp/tmpvmpbeshd.py", line 135, in _body
    ag__.if_stmt((ag__.ld(attention) is not None), if_body_5, else_body_5, get_state_5, set_state_5, ('attention_ta',), 1)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/operators/control_flow.py", line 1165, in if_stmt
    _py_if_stmt(cond, body, orelse)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/operators/control_flow.py", line 1218, in _py_if_stmt
    return body() if cond else orelse()
  File "/tmp/tmpvmpbeshd.py", line 130, in if_body_5
    attention_ta = ag__.converted_call(ag__.ld(attention_ta).write, (ag__.ld(step), ag__.ld(attention)), None, fscope_2)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 350, in converted_call
    return _call_unconverted(f, args, kwargs, options, False)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 479, in _call_unconverted
    return f(*args)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/util/tf_should_use.py", line 247, in wrapped
    return _add_should_use_warning(fn(*args, **kwargs),
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/tensor_array_ops.py", line 1159, in write
    return self._implementation.write(index, value, name=name)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/tensor_array_ops.py", line 536, in write
    _check_dtypes(value, self._dtype)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/tensor_array_ops.py", line 1315, in _check_dtypes
    "".join(traceback.format_stack())))

2021-05-25 12:51:22.729000: I control_flow.py:1218] Number of model parameters: 199372366
2021-05-25 12:51:23.987000: I control_flow.py:1218] Number of model weights: 42 (trainable = 42, non trainable = 0)
2021-05-25 12:51:24.138000: I coordinator.py:219] Error reported to Coordinator: in user code:

    /p/home/gerryc/OpenNMT-tf-latest/opennmt/training.py:370 _forward  *
        loss, gradients = self._compute_gradients(
    /p/home/gerryc/OpenNMT-tf-latest/opennmt/training.py:356 _compute_gradients  *
        gradients = self._optimizer.get_gradients(
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py:693 get_gradients  **
        grads = self._optimizer.get_gradients(loss, params)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:716 get_gradients
        grads = gradients.gradients(loss, params)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_impl.py:172 gradients
        unconnected_gradients)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:684 _GradientsHelper
        lambda: grad_fn(op, *out_grads))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:340 _MaybeCompile
        return grad_fn()  # Exit early
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:684 <lambda>
        lambda: grad_fn(op, *out_grads))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:358 _WhileGrad
        util.unique_grad_fn_name(body_graph.name), op, maximum_iterations)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:636 _create_grad_func
        body_graph_inputs, body_graph_outputs))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py:990 func_graph_from_py_func
        func_outputs = python_func(*func_args, **func_kwargs)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:632 <lambda>
        lambda *args: _grad_fn(ys, xs, args, body_graph),
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:692 _grad_fn
        unconnected_gradients="zero")
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:684 _GradientsHelper
        lambda: grad_fn(op, *out_grads))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:340 _MaybeCompile
        return grad_fn()  # Exit early
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:684 <lambda>
        lambda: grad_fn(op, *out_grads))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/list_ops.py:313 _TensorListSetItemGrad
        element_dtype=item.dtype)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/list_ops.py:111 tensor_list_get_item
        name=name)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gen_list_ops.py:565 tensor_list_get_item
        element_dtype=element_dtype, name=name)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/op_def_library.py:750 _apply_op_helper
        attrs=attr_protos, op_def=op_def)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:979 _create_op_internal
        compute_device=compute_device)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py:592 _create_op_internal
        compute_device)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/ops.py:3536 _create_op_internal
        op_def=op_def)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/ops.py:2016 __init__
        control_input_ops, op_def)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/ops.py:1856 _create_c_op
        raise ValueError(str(e))

    ValueError: Expected list with element dtype half but got list with element dtype float for '{{node Adam/gradients/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while_grad/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while/TensorArrayV2Write_1/TensorListSetItem_grad/TensorListGetItem}} = TensorListGetItem[element_dtype=DT_HALF](Adam/gradients/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while_grad/gradients/grad_ys_13, Adam/gradients/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while_grad/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while/TensorArrayV2Write/TensorListSetItem_grad/TensorListSetItem/TensorListPopBack:1, Adam/gradients/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while_grad/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while/TensorArrayV2Write_1/TensorListSetItem_grad/Shape)' with input shapes: [], [], [2].
Traceback (most recent call last):
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/training/coordinator.py", line 297, in stop_on_exception
    yield
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/distribute/mirrored_run.py", line 323, in run
    self.main_result = self.main_fn(*self.main_args, **self.main_kwargs)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/autograph/impl/api.py", line 670, in wrapper
    raise e.ag_error_metadata.to_exception(e)
ValueError: in user code:

    /p/home/gerryc/OpenNMT-tf-latest/opennmt/training.py:370 _forward  *
        loss, gradients = self._compute_gradients(
    /p/home/gerryc/OpenNMT-tf-latest/opennmt/training.py:356 _compute_gradients  *
        gradients = self._optimizer.get_gradients(
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py:693 get_gradients  **
        grads = self._optimizer.get_gradients(loss, params)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:716 get_gradients
        grads = gradients.gradients(loss, params)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_impl.py:172 gradients
        unconnected_gradients)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:684 _GradientsHelper
        lambda: grad_fn(op, *out_grads))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:340 _MaybeCompile
        return grad_fn()  # Exit early
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:684 <lambda>
        lambda: grad_fn(op, *out_grads))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:358 _WhileGrad
        util.unique_grad_fn_name(body_graph.name), op, maximum_iterations)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:636 _create_grad_func
        body_graph_inputs, body_graph_outputs))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py:990 func_graph_from_py_func
        func_outputs = python_func(*func_args, **func_kwargs)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:632 <lambda>
        lambda *args: _grad_fn(ys, xs, args, body_graph),
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:692 _grad_fn
        unconnected_gradients="zero")
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:684 _GradientsHelper
        lambda: grad_fn(op, *out_grads))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:340 _MaybeCompile
        return grad_fn()  # Exit early
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:684 <lambda>
        lambda: grad_fn(op, *out_grads))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/list_ops.py:313 _TensorListSetItemGrad
        element_dtype=item.dtype)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/list_ops.py:111 tensor_list_get_item
        name=name)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gen_list_ops.py:565 tensor_list_get_item
        element_dtype=element_dtype, name=name)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/op_def_library.py:750 _apply_op_helper
        attrs=attr_protos, op_def=op_def)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:979 _create_op_internal
        compute_device=compute_device)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py:592 _create_op_internal
        compute_device)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/ops.py:3536 _create_op_internal
        op_def=op_def)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/ops.py:2016 __init__
        control_input_ops, op_def)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/ops.py:1856 _create_c_op
        raise ValueError(str(e))

    ValueError: Expected list with element dtype half but got list with element dtype float for '{{node Adam/gradients/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while_grad/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while/TensorArrayV2Write_1/TensorListSetItem_grad/TensorListGetItem}} = TensorListGetItem[element_dtype=DT_HALF](Adam/gradients/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while_grad/gradients/grad_ys_13, Adam/gradients/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while_grad/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while/TensorArrayV2Write/TensorListSetItem_grad/TensorListSetItem/TensorListPopBack:1, Adam/gradients/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while_grad/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while/TensorArrayV2Write_1/TensorListSetItem_grad/Shape)' with input shapes: [], [], [2].

Traceback (most recent call last):
  File "opennmt/bin/main.py", line 363, in <module>
    main()
  File "opennmt/bin/main.py", line 326, in main
    hvd=hvd,
  File "/p/home/gerryc/OpenNMT-tf-latest/opennmt/runner.py", line 281, in train
    moving_average_decay=train_config.get("moving_average_decay"),
  File "/p/home/gerryc/OpenNMT-tf-latest/opennmt/training.py", line 123, in __call__
    dataset, accum_steps=accum_steps, report_steps=report_steps
  File "/p/home/gerryc/OpenNMT-tf-latest/opennmt/training.py", line 260, in _steps
    loss = forward_fn()
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 828, in __call__
    result = self._call(*args, **kwds)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 871, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 726, in _initialize
    *args, **kwds))
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 2969, in _get_concrete_function_internal_garbage_collected
    graph_function, _ = self._maybe_define_function(args, kwargs)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 3361, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 3206, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py", line 990, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 634, in wrapped_fn
    out = weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "/p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py", line 977, in wrapper
    raise e.ag_error_metadata.to_exception(e)
ValueError: in user code:

    /p/home/gerryc/OpenNMT-tf-latest/opennmt/training.py:245 _forward  *
        target,
    /p/home/gerryc/OpenNMT-tf-latest/opennmt/training.py:485 _forward  *
        per_replica_loss = self._strategy.run(
    /p/home/gerryc/OpenNMT-tf-latest/opennmt/training.py:370 _forward  *
        loss, gradients = self._compute_gradients(
    /p/home/gerryc/OpenNMT-tf-latest/opennmt/training.py:356 _compute_gradients  *
        gradients = self._optimizer.get_gradients(
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/keras/mixed_precision/loss_scale_optimizer.py:693 get_gradients  **
        grads = self._optimizer.get_gradients(loss, params)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/keras/optimizer_v2/optimizer_v2.py:716 get_gradients
        grads = gradients.gradients(loss, params)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_impl.py:172 gradients
        unconnected_gradients)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:684 _GradientsHelper
        lambda: grad_fn(op, *out_grads))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:340 _MaybeCompile
        return grad_fn()  # Exit early
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:684 <lambda>
        lambda: grad_fn(op, *out_grads))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:358 _WhileGrad
        util.unique_grad_fn_name(body_graph.name), op, maximum_iterations)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:636 _create_grad_func
        body_graph_inputs, body_graph_outputs))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py:990 func_graph_from_py_func
        func_outputs = python_func(*func_args, **func_kwargs)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:632 <lambda>
        lambda *args: _grad_fn(ys, xs, args, body_graph),
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:692 _grad_fn
        unconnected_gradients="zero")
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:684 _GradientsHelper
        lambda: grad_fn(op, *out_grads))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:340 _MaybeCompile
        return grad_fn()  # Exit early
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gradients_util.py:684 <lambda>
        lambda: grad_fn(op, *out_grads))
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/list_ops.py:313 _TensorListSetItemGrad
        element_dtype=item.dtype)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/list_ops.py:111 tensor_list_get_item
        name=name)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/gen_list_ops.py:565 tensor_list_get_item
        element_dtype=element_dtype, name=name)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/op_def_library.py:750 _apply_op_helper
        attrs=attr_protos, op_def=op_def)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/ops/while_v2.py:979 _create_op_internal
        compute_device=compute_device)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/func_graph.py:592 _create_op_internal
        compute_device)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/ops.py:3536 _create_op_internal
        op_def=op_def)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/ops.py:2016 __init__
        control_input_ops, op_def)
    /p/home/gerryc/.conda/envs/tf_venv/lib/python3.6/site-packages/tensorflow/python/framework/ops.py:1856 _create_c_op
        raise ValueError(str(e))

    ValueError: Expected list with element dtype half but got list with element dtype float for '{{node Adam/gradients/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while_grad/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while/TensorArrayV2Write_1/TensorListSetItem_grad/TensorListGetItem}} = TensorListGetItem[element_dtype=DT_HALF](Adam/gradients/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while_grad/gradients/grad_ys_13, Adam/gradients/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while_grad/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while/TensorArrayV2Write/TensorListSetItem_grad/TensorListSetItem/TensorListPopBack:1, Adam/gradients/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while_grad/gradients/nmt_big_v1_1/attentional_rnn_decoder_1/while/TensorArrayV2Write_1/TensorListSetItem_grad/Shape)' with input shapes: [], [], [2].

```






Discussion:

Resolving patch:
diff --git a/opennmt/decoders/decoder.py b/opennmt/decoders/decoder.py
--- a/opennmt/decoders/decoder.py
+++ b/opennmt/decoders/decoder.py
@@ -341,7 +341,7 @@ def _body(step, state, inputs, outputs_ta, attention_ta):
 
         step = tf.constant(0, dtype=tf.int32)
         outputs_ta = tf.TensorArray(inputs.dtype, size=max_step)
-        attention_ta = tf.TensorArray(tf.float32, size=max_step)
+        attention_ta = tf.TensorArray(inputs.dtype, size=max_step)
 
         _, state, _, outputs_ta, attention_ta = tf.while_loop(
             lambda *arg: True,


Test patch:
diff --git a/opennmt/tests/model_test.py b/opennmt/tests/model_test.py
--- a/opennmt/tests/model_test.py
+++ b/opennmt/tests/model_test.py
@@ -382,6 +382,19 @@ def testSequenceToSequenceServing(self):
         op_types = set(op.type for op in concrete_function.graph.get_operations())
         self.assertNotIn("Addons>GatherTree", op_types)
 
+    @test_util.run_with_mixed_precision
+    def testRNNWithMixedPrecision(self):
+        features_file, labels_file, data_config = self._makeToyEnDeData()
+        model = models.LuongAttention()
+        model.initialize(data_config)
+        dataset = model.examples_inputter.make_training_dataset(
+            features_file, labels_file, 16
+        )
+        features, labels = next(iter(dataset))
+        outputs, _ = model(features, labels=labels, training=True)
+        self.assertEqual(outputs["logits"].dtype, tf.float16)
+        self.assertEqual(outputs["attention"].dtype, tf.float16)
+
     @parameterized.expand(
         [
             [tf.estimator.ModeKeys.TRAIN],