additive attention layer a k a bahdanaustyle attention inputs are a list with 2 or 3 elements 1 a query tensor of shape batchsize tq dim 2 a value tensor of shape batchsize tv dim 3 a optional key tensor of shape batchsize tv dim if none supplied value will be used as key the calculation follows the steps 1 calculate attention scores using query and key with shape batchsize tq tv as a nonlinear sum scores reducesumtanhquery key axis1 2 use scores to calculate a softmax distribution with shape batchsize tq tv 3 use the softmax distribution to create a linear combination of value with shape batchsize tq dim args usescale if true will create a scalar variable to scale the attention scores dropout float between 0 and 1 fraction of the units to drop for the attention scores defaults to 0 0 call args inputs list of the following tensors query query tensor of shape batchsize tq dim value value tensor of shape batchsize tv dim key optional key tensor of shape batchsize tv dim if not given will use value for both key and value which is the most common case mask list of the following tensors querymask a boolean mask tensor of shape batchsize tq if given the output will be zero at the positions where maskfalse valuemask a boolean mask tensor of shape batchsize tv if given will apply the mask such that values at positions where maskfalse do not contribute to the result returnattentionscores bool it true returns the attention scores after masking and softmax as an additional output argument training python boolean indicating whether the layer should behave in training mode adding dropout or in inference mode no dropout usecausalmask boolean set to true for decoder selfattention adds a mask such that position i cannot attend to positions j i this prevents the flow of information from the future towards the past defaults to false output attention outputs of shape batchsize tq dim optional attention scores after masking and softmax with shape batchsize tq tv calculates attention scores as a nonlinear sum of query and key args query query tensor of shape batchsize tq dim key key tensor of shape batchsize tv dim returns tensor of shape batchsize tq tv reshape tensors to enable broadcasting reshape into batchsize tq 1 dim reshape into batchsize 1 tv dim additive attention layer a k a bahdanau style attention inputs are a list with 2 or 3 elements 1 a query tensor of shape batch_size tq dim 2 a value tensor of shape batch_size tv dim 3 a optional key tensor of shape batch_size tv dim if none supplied value will be used as key the calculation follows the steps 1 calculate attention scores using query and key with shape batch_size tq tv as a non linear sum scores reduce_sum tanh query key axis 1 2 use scores to calculate a softmax distribution with shape batch_size tq tv 3 use the softmax distribution to create a linear combination of value with shape batch_size tq dim args use_scale if true will create a scalar variable to scale the attention scores dropout float between 0 and 1 fraction of the units to drop for the attention scores defaults to 0 0 call args inputs list of the following tensors query query tensor of shape batch_size tq dim value value tensor of shape batch_size tv dim key optional key tensor of shape batch_size tv dim if not given will use value for both key and value which is the most common case mask list of the following tensors query_mask a boolean mask tensor of shape batch_size tq if given the output will be zero at the positions where mask false value_mask a boolean mask tensor of shape batch_size tv if given will apply the mask such that values at positions where mask false do not contribute to the result return_attention_scores bool it true returns the attention scores after masking and softmax as an additional output argument training python boolean indicating whether the layer should behave in training mode adding dropout or in inference mode no dropout use_causal_mask boolean set to true for decoder self attention adds a mask such that position i cannot attend to positions j i this prevents the flow of information from the future towards the past defaults to false output attention outputs of shape batch_size tq dim optional attention scores after masking and softmax with shape batch_size tq tv calculates attention scores as a nonlinear sum of query and key args query query tensor of shape batch_size tq dim key key tensor of shape batch_size tv dim returns tensor of shape batch_size tq tv reshape tensors to enable broadcasting reshape into batch_size tq 1 dim reshape into batch_size 1 tv dim
from keras import ops
from keras.api_export import keras_export
from keras.layers.attention.attention import Attention


@keras_export("keras.layers.AdditiveAttention")
class AdditiveAttention(Attention):
    

    def __init__(
        self,
        use_scale=True,
        dropout=0.0,
        **kwargs,
    ):
        super().__init__(use_scale=use_scale, dropout=dropout, **kwargs)

    def build(self, input_shape):
        self._validate_inputs(input_shape)
        dim = input_shape[0][-1]
        self.scale = None
        if self.use_scale:
            self.scale = self.add_weight(
                name="scale",
                shape=[dim],
                initializer="glorot_uniform",
                dtype=self.dtype,
                trainable=True,
            )
        self.built = True

    def _calculate_scores(self, query, key):
        
        
        
        q_reshaped = ops.expand_dims(query, axis=-2)
        
        k_reshaped = ops.expand_dims(key, axis=-3)
        scale = self.scale if self.use_scale else 1.0
        return ops.sum(scale * ops.tanh(q_reshaped + k_reshaped), axis=-1)

    def get_config(self):
        base_config = super().get_config()
        del base_config["score_mode"]
        return base_config