Cannot call AttentionLayer with symbolic tensors and return_attention_scores=True in Keras 3

[boweia]

In Keras 3 I cannot call the call method of AttentionLayer (or other attention layer classes) with a symbolic KerasTensor input and return_attention_scores=True.

I am using Python 3.12.1 and Keras 3.6.0

>>> in1 = keras.Input(shape=(10, 7))
>>> in2 = keras.Input(shape=(8, 7))
>>> attLayer = keras.layers.Attention()     
>>> out1, out2 = attLayer([in1, in2], return_attention_scores=True)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<venv_root>/lib/python3.12/site-packages/keras/src/backend/common/keras_tensor.py", line 167, in __iter__
    raise NotImplementedError(
NotImplementedError: Iterating over a symbolic KerasTensor is not supported.

The layer call method works when replacing the symbolic inputs with np.ndarray inputs.

[dhantule]

Hi @boweia, Thanks for reporting this.

You are trying to access attention scores with symbolic KerasTensor, which are only placeholders and do not hold any real data; while an eager tensor holds actual data and can be evaluated immediately in eager execution mode. Attention layer computes attention scores when the model is executed with real data.
If we pass the symbolic tensor to a Keras model, it will produce an eager tensor during the execution phase. Once the model is defined, called and data is passed to the model, you can get attention scores.

class AttentionModel(Model):
    def __init__(self):
        super(AttentionModel, self).__init__()
        self.attention = layers.Attention()

    def call(self,inputs):
        in1, in2 = inputs
        out1, out2 = self.attention([in1, in2], return_attention_scores=True)
        return out1, out2

in1 = keras.Input(shape=(10, 7))
in2 = keras.Input(shape=(8, 7))

model = AttentionModel()
out1, out2 =model([in1,in2]

[boweia]

Thanks @dhantule. I think what I'm unclear on is why the attention scores cannot be returned from the call method as a symbolic tensor, the same way the attention output is returned. Also, this workflow previously worked in TensorFlow 2.14 with Keras 2 and I only started seeing it after updating to newer versions recently.

If defining a wrapper class is the only way to support this I see that it can work, but it seems silly that its necessary instead of purely using the Keras functional API to define a model.

[dhantule]

Hi @boweia,
In TensorFlow 2.14 with Keras 2, eager execution is enabled by default. In Keras 3 with TensorFlow 2.14 and beyond while eager execution is still supported, the framework encourages graph execution (using tf.function) for better performance in production and scaling scenarios. In Keras 3, eager execution is no longer the default mode for all operations.
You can read more about migrating Keras 2 to Keras 3 here.

[Surya2k1]

Hi @boweia, In TensorFlow 2.14 with Keras 2, eager execution is enabled by default. In Keras 3 with TensorFlow 2.14 and beyond while eager execution is still supported, the framework encourages graph execution (using tf.function) for better performance in production and scaling scenarios. In Keras 3, eager execution is no longer the default mode for all operations. You can read more about migrating Keras 2 to Keras 3 here.

I don't think the issue is related to Graph or Eager execution. Calling layer.build along with layer.call will produce output and attention scores as Keras tensors.

in1 = keras.Input(shape=(10, 7))
in2 = keras.Input(shape=(8, 7))
attLayer = keras.layers.Attention()
attLayer.build([getattr(in1,'shape'),getattr(in2,'shape')])
out1,out2 = attLayer.call([in1, in2], return_attention_scores=True)

I believe there is an issue with layer.__call__ . It seems the kwargs being ignored in symbolic call. If someone confirms this needs fix may be I can look into it.

Or is there is any reason for current behaviour?