# Buggy code
def state_size(self):
    """The `state_size` property of `AttentionWrapper`.

    Returns:
      An `AttentionWrapperState` tuple containing shapes used
      by this object.
    """
    return AttentionWrapperState(
        cell_state=self._cell.state_size,
        time=tf.TensorShape([]),
        attention=self._get_attention_layer_size(),
        alignments=self._item_or_tuple(
            a.alignments_size for a in self._attention_mechanisms
        ),
        attention_state=self._item_or_tuple(
            a.state_size for a in self._attention_mechanisms
        ),
        alignment_history=self._item_or_tuple(
            a.alignments_size if self._alignment_history else ()
            for a in self._attention_mechanisms
        ),
    )  # sometimes a TensorArray


# 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 state_size(self):
    """The `state_size` property of `AttentionWrapper`.

    Returns:
      An `AttentionWrapperState` tuple containing shapes used
      by this object.
    """
    return AttentionWrapperState(
        cell_state=self._cell.state_size,
        attention=self._get_attention_layer_size(),
        alignments=self._item_or_tuple(
            a.alignments_size for a in self._attention_mechanisms
        ),
        attention_state=self._item_or_tuple(
            a.state_size for a in self._attention_mechanisms
        ),
        alignment_history=self._item_or_tuple(
            a.alignments_size if self._alignment_history else ()
            for a in self._attention_mechanisms
        ),
    )  # sometimes a TensorArray