Keras MultiHeadAttention layer: Dimension value must be integer or None or have an __index__ method, got value 'TensorShape([None, 6, 8])'

Question

I am running into a lot of trouble trying to use Keras MultiHeadAttention. My model:

def create_RNN_with_attention(input_shape, query):
    x=tf.keras.Input(shape=input_shape)
    Bdir_layer = Bidirectional(layers.LSTM(4, return_sequences=True, activation="tanh"))(x)
    Bdir_layer2 = Bidirectional(layers.LSTM(4, return_sequences=True, activation="tanh"))(Bdir_layer)

    attnLayer = MultiHeadAttention(num_heads=1, key_dim=5, value_dim=Bdir_layer2.shape)

    query_tensor = tf.keras.Input(shape=query.shape)
    value_tensor = tf.keras.Input(shape=Bdir_layer2.shape)
    output_attn = attnLayer(query_tensor, value_tensor)(Bdir_layer)(query)
    print(output_tensor.shape)
    
    outputs=Dense(1, trainable=True, activation="sigmoid")(output_attn)
    model=Model(x,outputs)    
    optm = keras.optimizers.Adam(learning_rate=0.005)
    model.compile(optimizer='adam',
                   loss = tf.keras.losses.BinaryCrossentropy(from_logits=False),
                   metrics=['accuracy'])    
    return model    


model_attention = create_RNN_with_attention(
                      input_shape=(X_train.shape[1], X_train.shape[2]),
                                    query=QUERY_train
                                  )
model_attention.summary()

My X_train has shape (22500, 6, 3), my y_train has shape (22500,) and my QUERY_train has shape ((22500, 5).

I get the exception when the following line is executed:

output_attn = attnLayer(query_tensor, value_tensor)(Bdir_layer)(query)

Exception:

Dimension value must be integer or None or have an __ index __ method, got value 'TensorShape([None, 6, 8])' with type '<class 'tensorflow.python.framework.tensor_shape.TensorShape'>'

Call arguments received: • query=tf.Tensor(shape=(None, 22500, 5), dtype=float32) • value=tf.Tensor(shape=(None, None, 6, 8), dtype=float32) • key=None • attention_mask=None • return_attention_scores=False • training=False

I figure there are probably more things wrong with this code... I am trying to implement this model: