How to overcome "TypeError: Exception encountered when calling layer "tf.keras.backend.rnn" (type TFOpLambda)"?

Question

I'm trying to re-implement the text summarization tutorial here. The tutorial employs the Attention Layer proposed for Keras on GitHub (which does not come with Keras). Getting the following error when I employ the Attention Layer:

/usr/local/lib/python3.7/dist-packages/keras/engine/keras_tensor.py in __array__(self, dtype)
    253   def __array__(self, dtype=None):
    254     raise TypeError(
--> 255         f'You are passing {self}, an intermediate Keras symbolic input/output, '
    256         'to a TF API that does not allow registering custom dispatchers, such '
    257         'as `tf.cond`, `tf.function`, gradient tapes, or `tf.map_fn`. '

TypeError: Exception encountered when calling layer "tf.keras.backend.rnn" (type TFOpLambda).

You are passing KerasTensor(type_spec=TensorSpec(shape=(None, 101), dtype=tf.float32, name=None), name='tf.compat.v1.nn.softmax_1/Softmax:0', description="created by layer 'tf.compat.v1.nn.softmax_1'"), an intermediate Keras symbolic input/output, to a TF API that does not allow registering custom dispatchers, such as `tf.cond`, `tf.function`, gradient tapes, or `tf.map_fn`. Keras Functional model construction only supports TF API calls that *do* support dispatching, such as `tf.math.add` or `tf.reshape`. Other APIs cannot be called directly on symbolic Kerasinputs/outputs. You can work around this limitation by putting the operation in a custom Keras layer `call` and calling that layer on this symbolic input/output.

Call arguments received:
  • step_function=<function AttentionLayer.call.<locals>.energy_step at 0x7f1d5ff279e0>
  • inputs=tf.Tensor(shape=(None, None, 256), dtype=float32)
  • initial_states=['tf.Tensor(shape=(None, 101), dtype=float32)']
  • go_backwards=False
  • mask=None
  • constants=None
  • unroll=False
  • input_length=None
  • time_major=False
  • zero_output_for_mask=False

How can I overcome this error? I've added my software stack below:

TensorFlow: `2.8.0`
Keras: `2.8.0`
Python Version: `3.7.12 (default, Jan 15 2022, 18:48:18)`