# Buggy code
def call(self, inputs, state, **kwargs):
    """Perform a step of attention-wrapped RNN.

    - Step 1: Mix the `inputs` and previous step's `attention` output via
      `cell_input_fn`.
    - Step 2: Call the wrapped `cell` with this input and its previous
      state.
    - Step 3: Score the cell's output with `attention_mechanism`.
    - Step 4: Calculate the alignments by passing the score through the
      `normalizer`.
    - Step 5: Calculate the context vector as the inner product between the
      alignments and the attention_mechanism's values (memory).
    - Step 6: Calculate the attention output by concatenating the cell
      output and context through the attention layer (a linear layer with
      `attention_layer_size` outputs).

    Args:
      inputs: (Possibly nested tuple of) Tensor, the input at this time
        step.
      state: An instance of `AttentionWrapperState` containing
        tensors from the previous time step.
      **kwargs: Dict, other keyword arguments for the cell call method.

    Returns:
      A tuple `(attention_or_cell_output, next_state)`, where:

      - `attention_or_cell_output` depending on `output_attention`.
      - `next_state` is an instance of `AttentionWrapperState`
         containing the state calculated at this time step.

    Raises:
      TypeError: If `state` is not an instance of `AttentionWrapperState`.
    """
    if not isinstance(state, AttentionWrapperState):
        try:
            state = AttentionWrapperState(*state)
        except TypeError:
            raise TypeError(
                "Expected state to be instance of AttentionWrapperState or "
                "values that can construct AttentionWrapperState. "
                "Received type %s instead." % type(state)
            )

    # Step 1: Calculate the true inputs to the cell based on the
    # previous attention value.
    cell_inputs = self._cell_input_fn(inputs, state.attention)
    cell_state = state.cell_state
    cell_output, next_cell_state = self._cell(cell_inputs, cell_state, **kwargs)

    cell_batch_size = (
        tf.compat.dimension_value(cell_output.shape[0]) or tf.shape(cell_output)[0]
    )
    error_message = (
        "When applying AttentionWrapper %s: "
        % self.name
        + "Non-matching batch sizes between the memory "
        "(encoder output) and the query (decoder output).  Are you using "
        "the BeamSearchDecoder?  You may need to tile your memory input "
        "via the tfa.seq2seq.tile_batch function with argument "
        "multiple=beam_width."
    )
    with tf.control_dependencies(
        self._batch_size_checks(cell_batch_size, error_message)
    ):  # pylint: disable=bad-continuation
        cell_output = tf.identity(cell_output, name="checked_cell_output")

    if self._is_multi:
        previous_attention_state = state.attention_state
        previous_alignment_history = state.alignment_history
    else:
        previous_attention_state = [state.attention_state]
        previous_alignment_history = [state.alignment_history]

    all_alignments = []
    all_attentions = []
    all_attention_states = []
    maybe_all_histories = []
    for i, attention_mechanism in enumerate(self._attention_mechanisms):
        attention, alignments, next_attention_state = self._attention_fn(
            attention_mechanism,
            cell_output,
            previous_attention_state[i],
            self._attention_layers[i] if self._attention_layers else None,
        )
        alignment_history = (
            previous_alignment_history[i].write(state.time, alignments)
            if self._alignment_history
            else ()
        )

        all_attention_states.append(next_attention_state)
        all_alignments.append(alignments)
        all_attentions.append(attention)
        maybe_all_histories.append(alignment_history)

    attention = tf.concat(all_attentions, 1)
    next_state = AttentionWrapperState(
        time=state.time + 1,
        cell_state=next_cell_state,
        attention=attention,
        attention_state=self._item_or_tuple(all_attention_states),
        alignments=self._item_or_tuple(all_alignments),
        alignment_history=self._item_or_tuple(maybe_all_histories),
    )

    if self._output_attention:
        return attention, next_state
    else:
        return cell_output, next_state


# Traceback (ValueError)
Test 1 passed
Traceback (most recent call last):
File "/run/media/zenbook/work/phd/work/__/pb1.py", line 26, in <module>
test(masked=True); print('Test 2 passed')
File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 580, in __call__
result = self._call(*args, **kwds)
File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 611, in _call
return self._stateless_fn(*args, **kwds)  # pylint: disable=not-callable
File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2419, in __call__
graph_function, args, kwargs = self._maybe_define_function(args, kwargs)
File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2777, in _maybe_define_function
graph_function = self._create_graph_function(args, kwargs)
File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2667, in _create_graph_function
capture_by_value=self._capture_by_value),
File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/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 "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 441, in wrapped_fn
return weak_wrapped_fn().__wrapped__(*args, **kwds)
File "/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py", line 968, in wrapper
raise e.ag_error_metadata.to_exception(e)
ValueError: in user code:

/run/media/zenbook/work/phd/work/__/pb1.py:22 test  *
return layer(data, mask=mask)
/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/keras/layers/recurrent.py:652 __call__  **
return super(RNN, self).__call__(inputs, **kwargs)
/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/keras/engine/base_layer.py:926 __call__
outputs = call_fn(cast_inputs, *args, **kwargs)
/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/keras/layers/recurrent.py:793 call
zero_output_for_mask=self.zero_output_for_mask)
/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/keras/backend.py:4205 rnn
**while_loop_kwargs)
/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/ops/control_flow_ops.py:2688 while_loop
back_prop=back_prop)
/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/ops/while_v2.py:229 while_loop
len_orig_loop_vars], expand_composites=True))
/home/zenbook/miniconda3/envs/tfn/lib/python3.7/site-packages/tensorflow/python/ops/while_v2.py:1143 _check_shapes_compat
"specify a less-specific shape." % (input_t.name, shape, t.shape))

ValueError: Input tensor 'rnn/AttentionWrapperZeroState/zeros_3:0' enters the loop with shape (), but has shape (4, 1) after one iteration. To allow the shape to vary across iterations, use the `shape_invariants` argument of tf.while_loop to specify a less-specific shape.


Process finished with exit code 1

# Fixed code
def call(self, inputs, state, **kwargs):
    """Perform a step of attention-wrapped RNN.

    - Step 1: Mix the `inputs` and previous step's `attention` output via
      `cell_input_fn`.
    - Step 2: Call the wrapped `cell` with this input and its previous
      state.
    - Step 3: Score the cell's output with `attention_mechanism`.
    - Step 4: Calculate the alignments by passing the score through the
      `normalizer`.
    - Step 5: Calculate the context vector as the inner product between the
      alignments and the attention_mechanism's values (memory).
    - Step 6: Calculate the attention output by concatenating the cell
      output and context through the attention layer (a linear layer with
      `attention_layer_size` outputs).

    Args:
      inputs: (Possibly nested tuple of) Tensor, the input at this time
        step.
      state: An instance of `AttentionWrapperState` containing
        tensors from the previous time step.
      **kwargs: Dict, other keyword arguments for the cell call method.

    Returns:
      A tuple `(attention_or_cell_output, next_state)`, where:

      - `attention_or_cell_output` depending on `output_attention`.
      - `next_state` is an instance of `AttentionWrapperState`
         containing the state calculated at this time step.

    Raises:
      TypeError: If `state` is not an instance of `AttentionWrapperState`.
    """
    if not isinstance(state, AttentionWrapperState):
        try:
            state = AttentionWrapperState(*state)
        except TypeError:
            raise TypeError(
                "Expected state to be instance of AttentionWrapperState or "
                "values that can construct AttentionWrapperState. "
                "Received type %s instead." % type(state)
            )

    # Step 1: Calculate the true inputs to the cell based on the
    # previous attention value.
    cell_inputs = self._cell_input_fn(inputs, state.attention)
    cell_state = state.cell_state
    cell_output, next_cell_state = self._cell(cell_inputs, cell_state, **kwargs)

    cell_batch_size = (
        tf.compat.dimension_value(cell_output.shape[0]) or tf.shape(cell_output)[0]
    )
    error_message = (
        "When applying AttentionWrapper %s: "
        % self.name
        + "Non-matching batch sizes between the memory "
        "(encoder output) and the query (decoder output).  Are you using "
        "the BeamSearchDecoder?  You may need to tile your memory input "
        "via the tfa.seq2seq.tile_batch function with argument "
        "multiple=beam_width."
    )
    with tf.control_dependencies(
        self._batch_size_checks(cell_batch_size, error_message)
    ):  # pylint: disable=bad-continuation
        cell_output = tf.identity(cell_output, name="checked_cell_output")

    if self._is_multi:
        previous_attention_state = state.attention_state
        previous_alignment_history = state.alignment_history
    else:
        previous_attention_state = [state.attention_state]
        previous_alignment_history = [state.alignment_history]

    all_alignments = []
    all_attentions = []
    all_attention_states = []
    maybe_all_histories = []
    for i, attention_mechanism in enumerate(self._attention_mechanisms):
        attention, alignments, next_attention_state = self._attention_fn(
            attention_mechanism,
            cell_output,
            previous_attention_state[i],
            self._attention_layers[i] if self._attention_layers else None,
        )
        alignment_history = (
            previous_alignment_history[i].write(
                previous_alignment_history[i].size(), alignments
            )
            if self._alignment_history
            else ()
        )

        all_attention_states.append(next_attention_state)
        all_alignments.append(alignments)
        all_attentions.append(attention)
        maybe_all_histories.append(alignment_history)

    attention = tf.concat(all_attentions, 1)
    next_state = AttentionWrapperState(
        cell_state=next_cell_state,
        attention=attention,
        attention_state=self._item_or_tuple(all_attention_states),
        alignments=self._item_or_tuple(all_alignments),
        alignment_history=self._item_or_tuple(maybe_all_histories),
    )

    if self._output_attention:
        return attention, next_state
    else:
        return cell_output, next_state