我想将tfa.seq2seq.BahdanauAttention与功能API tf.keras一起使用。我看了tensorflow/nmt/attention_model.py给出的示例。但是我不知道如何在tf.keras的功能性API中使用它。

因此，我想使用tfa.seq2seq.BahdanauAttention进行口述任务，如下所示：

    # Using tf.keras functional API
    encoder_out = a tensor of shape (batch_size,time_steps,units)
    decoder_out = a tensor of shape (batch_size,units)
    attn_out = attention_mechanism()(encoder_out,decoder_out)  # Need help figuring this out

谢谢。