# Buggy code
def call(self, inputs, memory=None, mask=None, cache=None, training=None):  # pylint: disable=arguments-differ
    """Runs the layer.

    Args:
      inputs: The sequence of queries. A tensor of shape :math:`[B, T_1, ...]`.
      memory: The sequence to attend. A tensor of shape :math:`[B, T_2, ...]`.
        If ``None``, computes self-attention.
      mask: The dot product mask. A boolean tensor of shape :math:`[B, T_2]` or
        :math:`[B, T_1, T_2]`.
      cache: An optional tuple containing projected keys and values from the
        previous step. Tensors of shape :math:`[B, H, T_2, D / H]`.
      training: Run in training mode.

    Returns:
      A tuple with the attention context, the updated cache and the attention
      probabilities of the first head (if :obj:`return_attention` is ``True``).
    """

    def _compute_kv(x):
        keys = self.linear_keys(x)
        keys = split_heads(keys, self.num_heads)
        values = self.linear_values(x)
        values = split_heads(values, self.num_heads)
        return keys, values

    # Compute queries.
    queries = self.linear_queries(inputs)
    queries = split_heads(queries, self.num_heads)
    queries *= self.num_units_per_head**-0.5

    # Compute keys and values.
    if memory is None:
        keys, values = _compute_kv(inputs)
        if cache:
            keys = tf.concat([cache[0], keys], axis=2)
            values = tf.concat([cache[1], values], axis=2)
    else:
        if cache:
            keys, values = tf.cond(
                tf.equal(tf.shape(cache[0])[2], 0),
                true_fn=lambda: _compute_kv(memory),
                false_fn=lambda: cache,
            )
        else:
            keys, values = _compute_kv(memory)

    if self.maximum_relative_position is not None:
        if memory is not None:
            raise ValueError(
                "Relative position representations only supports self-attention"
            )
        keys_length = tf.shape(keys)[2]
        relative_pos = relative_positions(
            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)
    else:
        relative_repr_keys = None
        relative_repr_values = None

    cache = (keys, values)

    # Dot product attention.
    dot = tf.matmul(queries, keys, transpose_b=True)
    if relative_repr_keys is not None:
        dot += matmul_with_relative_representations(
            queries, relative_repr_keys, transpose_b=True
        )
    if mask is not None:
        mask = tf.cast(mask, tf.float32)
        if mask.shape.rank == 2:
            mask = tf.expand_dims(mask, 1)  # Broadcast on time dimension.
        mask = tf.expand_dims(mask, 1)  # Broadcast on head dimension.
        dot = tf.cast(
            tf.cast(dot, tf.float32) * mask + ((1.0 - mask) * tf.float32.min), dot.dtype
        )
    attn = tf.cast(tf.nn.softmax(tf.cast(dot, tf.float32)), dot.dtype)
    drop_attn = common.dropout(attn, self.dropout, training=training)
    heads = tf.matmul(drop_attn, values)
    if relative_repr_values is not None:
        heads += matmul_with_relative_representations(drop_attn, relative_repr_values)

    # Concatenate all heads output.
    combined = combine_heads(heads)
    outputs = self.linear_output(combined)
    if self.return_attention:
        return outputs, cache, attn
    return outputs, cache


# Traceback (ValueError)
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)

# Fixed code
def call(self, inputs, memory=None, mask=None, cache=None, training=None):  # pylint: disable=arguments-differ
    """Runs the layer.

    Args:
      inputs: The sequence of queries. A tensor of shape :math:`[B, T_1, ...]`.
      memory: The sequence to attend. A tensor of shape :math:`[B, T_2, ...]`.
        If ``None``, computes self-attention.
      mask: The dot product mask. A boolean tensor of shape :math:`[B, T_2]` or
        :math:`[B, T_1, T_2]`.
      cache: An optional tuple containing projected keys and values from the
        previous step. Tensors of shape :math:`[B, H, T_2, D / H]`.
      training: Run in training mode.

    Returns:
      A tuple with the attention context, the updated cache and the attention
      probabilities of the first head (if :obj:`return_attention` is ``True``).
    """

    def _compute_kv(x):
        keys = self.linear_keys(x)
        keys = split_heads(keys, self.num_heads)
        values = self.linear_values(x)
        values = split_heads(values, self.num_heads)
        return keys, values

    # Compute queries.
    queries = self.linear_queries(inputs)
    queries = split_heads(queries, self.num_heads)
    queries *= self.num_units_per_head**-0.5

    # Compute keys and values.
    if memory is None:
        keys, values = _compute_kv(inputs)
        if cache:
            keys = tf.concat([cache[0], keys], axis=2)
            values = tf.concat([cache[1], values], axis=2)
    else:
        if cache:
            keys, values = tf.cond(
                tf.equal(tf.shape(cache[0])[2], 0),
                true_fn=lambda: _compute_kv(memory),
                false_fn=lambda: cache,
            )
        else:
            keys, values = _compute_kv(memory)

    if self.maximum_relative_position is not None:
        if memory is not None:
            raise ValueError(
                "Relative position representations only supports self-attention"
            )
        keys_length = tf.shape(keys)[2]
        relative_pos = relative_positions(
            keys_length, self.maximum_relative_position, with_cache=bool(cache)
        )
        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

    cache = (keys, values)

    # Dot product attention.
    dot = tf.matmul(queries, keys, transpose_b=True)
    if relative_repr_keys is not None:
        dot += matmul_with_relative_representations(
            queries, relative_repr_keys, transpose_b=True
        )
    if mask is not None:
        mask = tf.cast(mask, tf.float32)
        if mask.shape.rank == 2:
            mask = tf.expand_dims(mask, 1)  # Broadcast on time dimension.
        mask = tf.expand_dims(mask, 1)  # Broadcast on head dimension.
        dot = tf.cast(
            tf.cast(dot, tf.float32) * mask + ((1.0 - mask) * tf.float32.min), dot.dtype
        )
    attn = tf.cast(tf.nn.softmax(tf.cast(dot, tf.float32)), dot.dtype)
    drop_attn = common.dropout(attn, self.dropout, training=training)
    heads = tf.matmul(drop_attn, values)
    if relative_repr_values is not None:
        heads += matmul_with_relative_representations(drop_attn, relative_repr_values)

    # Concatenate all heads output.
    combined = combine_heads(heads)
    outputs = self.linear_output(combined)
    if self.return_attention:
        return outputs, cache, attn
    return outputs, cache