Self attention and multi input model in tensorflow

Question

I am trying to use Tensorflow Functional API (https://www.tensorflow.org/guide/keras/functional) to define a multi input neural network and add self attention layer after my