dotproduct attention layer a k a luongstyle 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 a key the calculation follows the steps 1 calculate attention scores using query and key with shape batchsize tq tv 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 seed a python integer to use as random seed incase of dropout scoremode function to use to compute attention scores one of dot concat dot refers to the dot product between the query and key vectors concat refers to the hyperbolic tangent of the concatenation of the query and key vectors 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 querykey dot product 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 applies attention scores to the given value tensor to use this method in your attention layer follow the steps use query tensor of shape batchsize tq and key tensor of shape batchsize tv to calculate the attention scores pass scores and value tensors to this method the method applies scoresmask calculates attentiondistribution softmaxscores then returns matmulattentiondistribution value apply querymask and return the result args scores scores float tensor of shape batchsize tq tv value value tensor of shape batchsize tv dim scoresmask a boolean mask tensor of shape batchsize 1 tv or batchsize tq tv if given scores at positions where scoresmaskfalse do not contribute to the result it must contain at least one true value in each line along the last dimension training python boolean indicating whether the layer should behave in training mode adding dropout or in inference mode no dropout returns tensor of shape batchsize tq dim attention scores after masking and softmax with shape batchsize tq tv bias so padding positions do not contribute to attention distribution note 65504 is the max float16 value mask of shape batchsize 1 tv creates a lower triangular mask so position i cannot attend to positions ji this prevents the flow of information from the future into the past causalmaskshape 1 tq tv mask of shape batchsize tq 1 validates arguments of the call method classname self class name if not isinstanceinputs list raise valueerror fclassname layer must be called on a list of inputs namely query value or query value key freceived inputsinputs if leninputs 2 or leninputs 3 raise valueerror fclassname layer accepts inputs list of length 2 or 3 namely query value or query value key freceived length leninputs if mask is not none if not isinstancemask list raise valueerror fclassname layer mask must be a list fnamely querymask valuemask received maskmask if lenmask 2 or lenmask 3 raise valueerror fclassname layer accepts mask list of length 2 or 3 freceived inputsinputs maskmask def getconfigself baseconfig super getconfig config usescale self usescale scoremode self scoremode dropout self dropout return baseconfig config dot product attention layer a k a luong 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 a key the calculation follows the steps 1 calculate attention scores using query and key with shape batch_size tq tv 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 seed a python integer to use as random seed incase of dropout score_mode function to use to compute attention scores one of dot concat dot refers to the dot product between the query and key vectors concat refers to the hyperbolic tangent of the concatenation of the query and key vectors 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 query key dot product 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 applies attention scores to the given value tensor to use this method in your attention layer follow the steps use query tensor of shape batch_size tq and key tensor of shape batch_size tv to calculate the attention scores pass scores and value tensors to this method the method applies scores_mask calculates attention_distribution softmax scores then returns matmul attention_distribution value apply query_mask and return the result args scores scores float tensor of shape batch_size tq tv value value tensor of shape batch_size tv dim scores_mask a boolean mask tensor of shape batch_size 1 tv or batch_size tq tv if given scores at positions where scores_mask false do not contribute to the result it must contain at least one true value in each line along the last dimension training python boolean indicating whether the layer should behave in training mode adding dropout or in inference mode no dropout returns tensor of shape batch_size tq dim attention scores after masking and softmax with shape batch_size tq tv bias so padding positions do not contribute to attention distribution note 65504 is the max float16 value mask of shape batch_size 1 tv creates a lower triangular mask so position i cannot attend to positions j i this prevents the flow of information from the future into the past causal_mask_shape 1 tq tv mask of shape batch_size tq 1 validates arguments of the call method
from keras import backend
from keras import ops
from keras.api_export import keras_export
from keras.layers.layer import Layer


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

    def __init__(
        self,
        use_scale=False,
        score_mode="dot",
        dropout=0.0,
        seed=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.use_scale = use_scale
        self.score_mode = score_mode
        self.dropout = dropout
        if self.dropout > 0:
            self.seed_generator = backend.random.SeedGenerator(seed=seed)

        if self.score_mode not in ["dot", "concat"]:
            raise ValueError(
                "Invalid value for argument score_mode. "
                "Expected one of {'dot', 'concat'}. "
                f"Received: score_mode={score_mode}"
            )

    def build(self, input_shape):
        self._validate_inputs(input_shape)
        self.scale = None
        self.concat_score_weight = None
        if self.use_scale:
            self.scale = self.add_weight(
                name="scale",
                shape=(),
                initializer="ones",
                dtype=self.dtype,
                trainable=True,
            )
        if self.score_mode == "concat":
            self.concat_score_weight = self.add_weight(
                name="concat_score_weight",
                shape=(),
                initializer="ones",
                dtype=self.dtype,
                trainable=True,
            )
        self.built = True

    def _calculate_scores(self, query, key):
        
        if self.score_mode == "dot":
            scores = ops.matmul(query, ops.transpose(key, axes=[0, 2, 1]))
            if self.scale is not None:
                scores *= self.scale
        elif self.score_mode == "concat":
            
            
            q_reshaped = ops.expand_dims(query, axis=-2)
            
            k_reshaped = ops.expand_dims(key, axis=-3)
            if self.scale is not None:
                scores = self.concat_score_weight * ops.sum(
                    ops.tanh(self.scale * (q_reshaped + k_reshaped)), axis=-1
                )
            else:
                scores = self.concat_score_weight * ops.sum(
                    ops.tanh(q_reshaped + k_reshaped), axis=-1
                )

        return scores

    def _apply_scores(self, scores, value, scores_mask=None, training=False):
        
        if scores_mask is not None:
            padding_mask = ops.logical_not(scores_mask)
            
            
            max_value = 65504.0 if scores.dtype == "float16" else 1.0e9
            scores -= max_value * ops.cast(padding_mask, dtype=scores.dtype)

        weights = ops.softmax(scores, axis=-1)
        if training and self.dropout > 0:
            weights = backend.random.dropout(
                weights,
                self.dropout,
                seed=self.seed_generator,
            )
        return ops.matmul(weights, value), weights

    def _calculate_score_mask(self, scores, v_mask, use_causal_mask):
        if v_mask is not None:
            
            v_mask = ops.expand_dims(v_mask, axis=-2)
        if not use_causal_mask:
            return v_mask

        
        
        
        score_shape = ops.shape(scores)
        
        mask_shape = (1, score_shape[-2], score_shape[-1])
        ones_mask = ops.ones(shape=mask_shape, dtype="int32")
        row_index = ops.cumsum(ones_mask, axis=-2)
        col_index = ops.cumsum(ones_mask, axis=-1)
        causal_mask = ops.greater_equal(row_index, col_index)

        if v_mask is not None:
            return causal_mask
        return ops.logical_and(v_mask, causal_mask)

    def call(
        self,
        inputs,
        mask=None,
        training=False,
        return_attention_scores=False,
        use_causal_mask=False,
    ):
        self._validate_inputs(inputs=inputs, mask=mask)
        q = inputs[0]
        v = inputs[1]
        k = inputs[2] if len(inputs) > 2 else v
        q_mask = mask[0] if mask else None
        v_mask = mask[1] if mask else None
        scores = self._calculate_scores(query=q, key=k)
        scores_mask = self._calculate_score_mask(
            scores, v_mask, use_causal_mask
        )
        result, attention_scores = self._apply_scores(
            scores=scores, value=v, scores_mask=scores_mask, training=training
        )
        if q_mask is not None:
            
            q_mask = ops.expand_dims(q_mask, axis=-1)
            result *= ops.cast(q_mask, dtype=result.dtype)
        if return_attention_scores:
            return result, attention_scores
        return result

    def compute_mask(self, inputs, mask=None):
        self._validate_inputs(inputs=inputs, mask=mask)
        if mask is None or mask[0] is None:
            return None
        return ops.convert_to_tensor(mask[0])

    def compute_output_shape(self, input_shape):
        return input_shape[0]

    def _validate_inputs(self, inputs, mask=None):
        
        class_name = self.__class__.__name__
        if not isinstance(inputs, list):
            raise ValueError(
                f"{class_name} layer must be called on a list of inputs, "
                "namely [query, value] or [query, value, key]. "
                f"Received: inputs={inputs}."
            )
        if len(inputs) < 2 or len(inputs) > 3:
            raise ValueError(
                f"{class_name} layer accepts inputs list of length 2 or 3, "
                "namely [query, value] or [query, value, key]. "
                f"Received length: {len(inputs)}."
            )
        if mask is not None:
            if not isinstance(mask, list):
                raise ValueError(
                    f"{class_name} layer mask must be a list, "
                    f"namely [query_mask, value_mask]. Received: mask={mask}."
                )
            if len(mask) < 2 or len(mask) > 3:
                raise ValueError(
                    f"{class_name} layer accepts mask list of length 2 or 3. "
                    f"Received: inputs={inputs}, mask={mask}."
                )

    def get_config(self):
        base_config = super().get_config()
        config = {
            "use_scale": self.use_scale,
            "score_mode": self.score_mode,
            "dropout": self.dropout,
        }
        return {**base_config, **config}