The attention scores are always `None` in `CachedMultiHeadAttention`

[apehex]

Describe the bug

The variable attention_scores introduced at line 111 is always None.

To Reproduce

Since it is an internal variable, I copied the subclass CMHA in this script:
https://colab.research.google.com/drive/1ZUS4mjDQktovKiJ8TQ7zYtm4PGjesXvG?usp=sharing

Expected behavior

The variable attention_scores should contain the cross correlation between query and key, which is useful for debugging a model IMHO.

Additional context

In recent Keras versions, the parent class MultiHeadAttention saves the argument return_attention_scores in self._return_attention_scores.

Then, the method _compute_attention checks this private property to decide whether or not to return the scores.
Since this state is not updated in CachedMultiHeadAttention.call, the attention scores will never be returned.

I'll also submit an issue to Keras to turn the attribute _return_attention_scores into an argument.

Would you like to help us fix it?

Yes, I have two potential fixes:

1. ignore attention scores entirely, which is consistent since the corresponding argument has been removed from CMHA
2. add the relevant argument and set the class property _return_attention_scores accordingly

WDYT?

[GargiGupta-io]

I’m planning to pick this up if no one is actively working on it.

I saw #2213 was closed after the added test failed and the PR went stale. Based on the review feedback there, the fix direction seems to be:

• add return_attention_scores as a call() argument, not an __init__() argument
• keep the return logic flat for the combinations of return_attention_scores and cache
• update the docstring return section
• add a unit test that verifies attention scores are not None and have the expected shape

Does that still sound like the right direction for this issue?