def test_return_attention_scores_true(self):
        """Test that the layer returns attention scores along with outputs."""
        # Generate dummy input data
        query = np.random.random((2, 8, 16)).astype(np.float32)
        value = np.random.random((2, 4, 16)).astype(np.float32)

        # Initialize the Attention layer
        layer = layers.Attention()

        # Call the layer with return_attention_scores=True
        output, attention_scores = layer(
            [query, value], return_attention_scores=True
        )

        # Check the shape of the outputs
        self.assertEqual(output.shape, (2, 8, 16))  # Output shape
        self.assertEqual(
            attention_scores.shape, (2, 8, 4)
        )  # Attention scores shape
Test that the layer returns attention scores along with outputs.
keras/src/layers/attention/attention_test.py