def test_query_mask_propagation(self):
        """Test automatic propagation of the query's mask."""
        try:
            layer = layers.MultiHeadAttention(num_heads=2, key_dim=2)
            self.assertTrue(layer.supports_masking)
            query = np.array(
                [[1, 2, 3, 0, 0], [3, 3, 1, 1, 2], [1, 0, 0, 0, 0]]
            )
            masked_query = layers.Embedding(4, 8, mask_zero=True)(query)
            query_mask = backend.get_keras_mask(masked_query)
            value = np.random.normal(size=(3, 3, 8))
            output = layer(query=masked_query, value=value)
        except RuntimeError as e:
            if e.args[0].startswith(
                "(*bias): last dimension must be contiguous"
            ):
                self.skipTest(
                    "PyTorch errors out on GPU: issue to track bug is here "
                    "https://github.com/keras-team/keras/issues/20459"
                )
        self.assertAllClose(query_mask, output._keras_mask)
Test automatic propagation of the query's mask.
keras/src/layers/attention/multi_head_attention_test.py