Repository: OpenNMT/OpenNMT-tf

Issue:
Crash when training TransformerBaseRelative model
When training a TransformerBaseRelative model using the following command:
```
onmt-main --model_type TransformerBaseRelative --config config.yml --auto_config train --with_eval
```
I receive the following error:
```
Traceback (most recent call last):
  File "C:\Users\damie\AppData\Local\Programs\Python\Python37\lib\runpy.py", line 193, in _run_module_as_main
    "__main__", mod_spec)
  File "C:\Users\damie\AppData\Local\Programs\Python\Python37\lib\runpy.py", line 85, in _run_code
    exec(code, run_globals)
  File "C:\Users\damie\AppData\Local\pypoetry\Cache\virtualenvs\nlp-mt-testing-U3-1H1OD-py3.7\Scripts\onmt-main.exe\__main__.py", line 9, in <module>
  File "c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\opennmt\bin\main.py", line 224, in main
    hvd=hvd)
  File "c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\opennmt\runner.py", line 236, in train
    moving_average_decay=train_config.get("moving_average_decay"))
  File "c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\opennmt\training.py", line 93, in __call__
    for loss in self._steps(dataset, accum_steps=accum_steps, report_steps=report_steps):
  File "c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\opennmt\training.py", line 223, in _steps
    _step()
  File "c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\tensorflow\python\eager\def_function.py", line 580, in __call__
    result = self._call(*args, **kwds)
  File "c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\tensorflow\python\eager\def_function.py", line 627, in _call
    self._initialize(args, kwds, add_initializers_to=initializers)
  File "c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\tensorflow\python\eager\def_function.py", line 506, in _initialize
    *args, **kwds))
  File "c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\tensorflow\python\eager\function.py", line 2446, in _get_concrete_function_internal_garbage_collected
    graph_function, _, _ = self._maybe_define_function(args, kwargs)
  File "c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\tensorflow\python\eager\function.py", line 2777, in _maybe_define_function
    graph_function = self._create_graph_function(args, kwargs)
  File "c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\tensorflow\python\eager\function.py", line 2667, in _create_graph_function
    capture_by_value=self._capture_by_value),
  File "c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\tensorflow\python\framework\func_graph.py", line 981, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\tensorflow\python\eager\def_function.py", line 441, in wrapped_fn
    return weak_wrapped_fn().__wrapped__(*args, **kwds)
  File "c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\tensorflow\python\framework\func_graph.py", line 968, in wrapper
    raise e.ag_error_metadata.to_exception(e)
ValueError: in user code:

    c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\opennmt\training.py:215 _step  *
        return self._step()
    c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\opennmt\training.py:314 _step  *
        self._gradient_accumulator.reset()
    c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\opennmt\optimizers\utils.py:124 reset  *
        gradient.assign(tf.zeros(gradient.shape, dtype=gradient.dtype))
    c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\tensorflow\python\ops\array_ops.py:2677 wrapped  **
        tensor = fun(*args, **kwargs)
    c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\tensorflow\python\ops\array_ops.py:2730 zeros
        shape = ops.convert_to_tensor(shape, dtype=dtypes.int32)
    c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\tensorflow\python\framework\ops.py:1341 convert_to_tensor
        ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
    c:\users\damie\appdata\local\pypoetry\cache\virtualenvs\nlp-mt-testing-u3-1h1od-py3.7\lib\site-packages\tensorflow\python\framework\constant_op.py:338 _tensor_shape_tensor_conversion_function
        "Cannot convert a partially known TensorShape to a Tensor: %s" % s)

    ValueError: Cannot convert a partially known TensorShape to a Tensor: (None, 64)
```
This only seems to occur in OpenNMT-tf 2.11. My config file does not set any hyperparameters. It only configures the location of data. Training using TransformerBase does not produce the error. I am running on Windows 10.


Discussion:

Resolving patch:
diff --git a/opennmt/layers/transformer.py b/opennmt/layers/transformer.py
--- a/opennmt/layers/transformer.py
+++ b/opennmt/layers/transformer.py
@@ -272,8 +272,8 @@ def _compute_kv(x):
           keys_length,
           self.maximum_relative_position,
           with_cache=bool(cache))
-      relative_repr_keys = tf.gather(self.relative_position_keys, relative_pos)
-      relative_repr_values = tf.gather(self.relative_position_values, relative_pos)
+      relative_repr_keys = tf.nn.embedding_lookup(self.relative_position_keys, relative_pos)
+      relative_repr_values = tf.nn.embedding_lookup(self.relative_position_values, relative_pos)
     else:
       relative_repr_keys = None
       relative_repr_values = None


Test patch:
diff --git a/opennmt/tests/transformer_test.py b/opennmt/tests/transformer_test.py
--- a/opennmt/tests/transformer_test.py
+++ b/opennmt/tests/transformer_test.py
@@ -142,6 +142,20 @@ def testMultiHeadSelfAttentionRelativePositionsWithCache(self):
     cache = (tf.zeros([4, 4, 0, 5]), tf.zeros([4, 4, 0, 5]))
     _, cache = attention(x, cache=cache)
 
+  def testMultiHeadSelfAttentionRelativeGradients(self):
+    attention = transformer.MultiHeadAttention(4, 20, maximum_relative_position=6)
+
+    @tf.function
+    def _compute_gradients_in_function(x):
+      with tf.GradientTape() as tape:
+        y, _ = attention(x)
+        loss = tf.math.reduce_sum(y)
+      gradients = tape.gradient(loss, attention.weights)
+      for gradient in gradients:
+        self.assertTrue(gradient.shape.is_fully_defined())
+
+    _compute_gradients_in_function(tf.random.uniform([4, 1, 10]))
+
   def testMultiHeadAttention(self):
     attention = transformer.MultiHeadAttention(4, 20)
     queries = tf.random.uniform([4, 5, 10])