# Buggy code
def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
    """Return an initial (zero) state tuple for this `AttentionWrapper`.

    **NOTE** Please see the initializer documentation for details of how
    to call `get_initial_state` if using an `AttentionWrapper` with a
    `BeamSearchDecoder`.

    Args:
      inputs: The inputs that will be fed to this cell.
      batch_size: `0D` integer tensor: the batch size.
      dtype: The internal state data type.

    Returns:
      An `AttentionWrapperState` tuple containing zeroed out tensors and,
      possibly, empty `TensorArray` objects.

    Raises:
      ValueError: (or, possibly at runtime, InvalidArgument), if
        `batch_size` does not match the output size of the encoder passed
        to the wrapper object at initialization time.
    """
    if inputs is not None:
        batch_size = tf.shape(inputs)[0]
        dtype = inputs.dtype
    with tf.name_scope(type(self).__name__ + "ZeroState"):  # pylint: disable=bad-continuation
        if self._initial_cell_state is not None:
            cell_state = self._initial_cell_state
        else:
            cell_state = self._cell.get_initial_state(
                batch_size=batch_size, dtype=dtype
            )
        error_message = (
            "When calling get_initial_state of AttentionWrapper %s: "
            % self.name
            + "Non-matching batch sizes between the memory "
            "(encoder output) and the requested batch size. Are you using "
            "the BeamSearchDecoder?  If so, make sure your encoder output "
            "has been tiled to beam_width via "
            "tfa.seq2seq.tile_batch, and the batch_size= argument "
            "passed to get_initial_state is batch_size * beam_width."
        )
        with tf.control_dependencies(
            self._batch_size_checks(batch_size, error_message)
        ):  # pylint: disable=bad-continuation
            cell_state = tf.nest.map_structure(
                lambda s: tf.identity(s, name="checked_cell_state"), cell_state
            )
        initial_alignments = [
            attention_mechanism.initial_alignments(batch_size, dtype)
            for attention_mechanism in self._attention_mechanisms
        ]
        return AttentionWrapperState(
            cell_state=cell_state,
            time=tf.zeros([], dtype=tf.int32),
            attention=tf.zeros(
                [batch_size, self._get_attention_layer_size()], dtype=dtype
            ),
            alignments=self._item_or_tuple(initial_alignments),
            attention_state=self._item_or_tuple(
                attention_mechanism.initial_state(batch_size, dtype)
                for attention_mechanism in self._attention_mechanisms
            ),
            alignment_history=self._item_or_tuple(
                tf.TensorArray(
                    dtype, size=0, dynamic_size=True, element_shape=alignment.shape
                )
                if self._alignment_history
                else ()
                for alignment in initial_alignments
            ),
        )


# 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 get_initial_state(self, inputs=None, batch_size=None, dtype=None):
    """Return an initial (zero) state tuple for this `AttentionWrapper`.

    **NOTE** Please see the initializer documentation for details of how
    to call `get_initial_state` if using an `AttentionWrapper` with a
    `BeamSearchDecoder`.

    Args:
      inputs: The inputs that will be fed to this cell.
      batch_size: `0D` integer tensor: the batch size.
      dtype: The internal state data type.

    Returns:
      An `AttentionWrapperState` tuple containing zeroed out tensors and,
      possibly, empty `TensorArray` objects.

    Raises:
      ValueError: (or, possibly at runtime, InvalidArgument), if
        `batch_size` does not match the output size of the encoder passed
        to the wrapper object at initialization time.
    """
    if inputs is not None:
        batch_size = tf.shape(inputs)[0]
        dtype = inputs.dtype
    with tf.name_scope(type(self).__name__ + "ZeroState"):  # pylint: disable=bad-continuation
        if self._initial_cell_state is not None:
            cell_state = self._initial_cell_state
        else:
            cell_state = self._cell.get_initial_state(
                batch_size=batch_size, dtype=dtype
            )
        error_message = (
            "When calling get_initial_state of AttentionWrapper %s: "
            % self.name
            + "Non-matching batch sizes between the memory "
            "(encoder output) and the requested batch size. Are you using "
            "the BeamSearchDecoder?  If so, make sure your encoder output "
            "has been tiled to beam_width via "
            "tfa.seq2seq.tile_batch, and the batch_size= argument "
            "passed to get_initial_state is batch_size * beam_width."
        )
        with tf.control_dependencies(
            self._batch_size_checks(batch_size, error_message)
        ):  # pylint: disable=bad-continuation
            cell_state = tf.nest.map_structure(
                lambda s: tf.identity(s, name="checked_cell_state"), cell_state
            )
        initial_alignments = [
            attention_mechanism.initial_alignments(batch_size, dtype)
            for attention_mechanism in self._attention_mechanisms
        ]
        return AttentionWrapperState(
            cell_state=cell_state,
            attention=tf.zeros(
                [batch_size, self._get_attention_layer_size()], dtype=dtype
            ),
            alignments=self._item_or_tuple(initial_alignments),
            attention_state=self._item_or_tuple(
                attention_mechanism.initial_state(batch_size, dtype)
                for attention_mechanism in self._attention_mechanisms
            ),
            alignment_history=self._item_or_tuple(
                tf.TensorArray(
                    dtype, size=0, dynamic_size=True, element_shape=alignment.shape
                )
                if self._alignment_history
                else ()
                for alignment in initial_alignments
            ),
        )