Since we are using TensorFlow for this lab, you can use the default one in the TF library: tf.keras.activations.softmax
Either import it (1) or use the fully qualified version (2):
(1) from tensorflow.keras.activations import softmax
(2) attention_weights = tf.keras.activations.softmax(…)