grouped query attention layer this is an implementation of groupedquery attention introduced by ainslie et al 2023https arxiv orgabs2305 13245 here numkeyvalueheads denotes number of groups setting numkeyvalueheads to 1 is equivalent to multiquery attention and when numkeyvalueheads is equal to numqueryheads it is equivalent to multihead attention this layer first projects query key and value tensors then key and value are repeated to match the number of heads of query then the query is scaled and dotproducted with key tensors these are softmaxed to obtain attention probabilities the value tensors are then interpolated by these probabilities and concatenated back to a single tensor args headdim size of each attention head numqueryheads number of query attention heads numkeyvalueheads number of key and value attention heads dropout dropout probability usebias boolean whether the dense layers use bias vectorsmatrices kernelinitializer initializer for dense layer kernels biasinitializer initializer for dense layer biases kernelregularizer regularizer for dense layer kernels biasregularizer regularizer for dense layer biases activityregularizer regularizer for dense layer activity kernelconstraint constraint for dense layer kernels biasconstraint constraint for dense layer kernels call arguments query query tensor of shape batchdim targetseqlen featuredim where batchdim is batch size targetseqlen is the length of target sequence and featuredim is dimension of feature value value tensor of shape batchdim sourceseqlen featuredim where batchdim is batch size sourceseqlen is the length of source sequence and featuredim is dimension of feature key optional key tensor of shape batchdim sourceseqlen featuredim if not given will use value for both key and value which is most common case attentionmask a boolean mask of shape batchdim targetseqlen sourceseqlen that prevents attention to certain positions the boolean mask specifies which query elements can attend to which key elements where 1 indicates attention and 0 indicates no attention broadcasting can happen for the missing batch dimensions and the head dimension returnattentionscores a boolean to indicate whether the output should be attentionoutput attentionscores if true or attentionoutput if false defaults to false training python boolean indicating whether the layer should behave in training mode adding dropout or in inference mode no dropout will go with either using the training mode of the parent layermodel or false inference if there is no parent layer usecausalmask a boolean to indicate whether to apply a causal mask to prevent tokens from attending to future tokens e g used in a decoder transformer returns attentionoutput result of the computation of shape batchdim targetseqlen featuredim where targetseqlen is for target sequence length and featuredim is the query input last dim attentionscores optional attention coefficients of shape batchdim numqueryheads targetseqlen sourceseqlen einsum variables b batch size q query length k keyvalue length m model dim u num query heads v num keyvalue heads h head dim create new clone of kernelbias initializer so that we don t reuse the initializer instance which could lead to same init value since initializer is stateless computes the attention mask using the keras masks of the inputs the query s mask is reshaped from b t to b t 1 the value s mask is reshaped from b s to b 1 s the key s mask is reshaped from b s to b 1 s the key s mask is ignored if key is none or if key is value if usecausalmasktrue then the causal mask is computed its shape is 1 t s all defined masks are merged using a logical and operation in general if the query and value are masked then there is no need to define the attentionmask args query projected query tensor of shape b t n keydim key projected key tensor of shape b t n keydim value projected value tensor of shape b t n valuedim attentionmask a boolean mask of shape b t s that prevents attention to certain positions usecausalmask a boolean to indicate whether to apply a causal mask to prevent tokens from attending to future tokens e g used in a decoder transformer returns attentionmask a boolean mask of shape b t s that prevents attention to certain positions based on the keras masks of the query key value and attentionmask tensors and the causal mask if usecausalmasktrue b batch size t max query length b batch size s max value length b batch size s max key length max value length the shape of the causal mask is 1 t s merge attentionmask automatic mask to shape b t s computes a causal mask e g for masked selfattention layers for example if query and value both contain sequences of length 4 this function returns a boolean tensor equal to true false false false true true false false true true true false true true true true args query query tensor of shape b t value value tensor of shape b s optional defaults to query returns mask a boolean tensor of shape 1 t s containing a lower triangular matrix of shape t s take the dot product between query and key to get the raw attention scores this is actually dropping out entire tokens to attend to which might seem a bit unusual but is taken from the original transformer paper normalize the attention scores to probabilities scores b n t s the expand dim happens starting from the numheads dimension batchdims numheads queryattentiondims keyattentiondims grouped query attention layer this is an implementation of grouped query attention introduced by ainslie et al 2023 https arxiv org abs 2305 13245 here num_key_value_heads denotes number of groups setting num_key_value_heads to 1 is equivalent to multi query attention and when num_key_value_heads is equal to num_query_heads it is equivalent to multi head attention this layer first projects query key and value tensors then key and value are repeated to match the number of heads of query then the query is scaled and dot producted with key tensors these are softmaxed to obtain attention probabilities the value tensors are then interpolated by these probabilities and concatenated back to a single tensor args head_dim size of each attention head num_query_heads number of query attention heads num_key_value_heads number of key and value attention heads dropout dropout probability use_bias boolean whether the dense layers use bias vectors matrices kernel_initializer initializer for dense layer kernels bias_initializer initializer for dense layer biases kernel_regularizer regularizer for dense layer kernels bias_regularizer regularizer for dense layer biases activity_regularizer regularizer for dense layer activity kernel_constraint constraint for dense layer kernels bias_constraint constraint for dense layer kernels call arguments query query tensor of shape batch_dim target_seq_len feature_dim where batch_dim is batch size target_seq_len is the length of target sequence and feature_dim is dimension of feature value value tensor of shape batch_dim source_seq_len feature_dim where batch_dim is batch size source_seq_len is the length of source sequence and feature_dim is dimension of feature key optional key tensor of shape batch_dim source_seq_len feature_dim if not given will use value for both key and value which is most common case attention_mask a boolean mask of shape batch_dim target_seq_len source_seq_len that prevents attention to certain positions the boolean mask specifies which query elements can attend to which key elements where 1 indicates attention and 0 indicates no attention broadcasting can happen for the missing batch dimensions and the head dimension return_attention_scores a boolean to indicate whether the output should be attention_output attention_scores if true or attention_output if false defaults to false training python boolean indicating whether the layer should behave in training mode adding dropout or in inference mode no dropout will go with either using the training mode of the parent layer model or false inference if there is no parent layer use_causal_mask a boolean to indicate whether to apply a causal mask to prevent tokens from attending to future tokens e g used in a decoder transformer returns attention_output result of the computation of shape batch_dim target_seq_len feature_dim where target_seq_len is for target sequence length and feature_dim is the query input last dim attention_scores optional attention coefficients of shape batch_dim num_query_heads target_seq_len source_seq_len einsum variables b batch size q query length k key value length m model dim u num query heads v num key value heads h head dim create new clone of kernel bias initializer so that we don t reuse the initializer instance which could lead to same init value since initializer is stateless batch_dim source_seq_len query_heads head_dim batch_dim source_seq_len query_heads head_dim batch_dim target_seq_len feature_dim computes the attention mask using the keras masks of the inputs the query s mask is reshaped from b t to b t 1 the value s mask is reshaped from b s to b 1 s the key s mask is reshaped from b s to b 1 s the key s mask is ignored if key is none or if key is value if use_causal_mask true then the causal mask is computed its shape is 1 t s all defined masks are merged using a logical and operation in general if the query and value are masked then there is no need to define the attention_mask args query projected query tensor of shape b t n key_dim key projected key tensor of shape b t n key_dim value projected value tensor of shape b t n value_dim attention_mask a boolean mask of shape b t s that prevents attention to certain positions use_causal_mask a boolean to indicate whether to apply a causal mask to prevent tokens from attending to future tokens e g used in a decoder transformer returns attention_mask a boolean mask of shape b t s that prevents attention to certain positions based on the keras masks of the query key value and attention_mask tensors and the causal mask if use_causal_mask true defensive casting b batch size t max query length shape is b t 1 defensive casting b batch size s max value length shape is b 1 s defensive casting b batch size s max key length max value length shape is b 1 s the shape of the causal mask is 1 t s merge attention_mask automatic mask to shape b t s computes a causal mask e g for masked self attention layers for example if query and value both contain sequences of length 4 this function returns a boolean tensor equal to true false false false true true false false true true true false true true true true args query query tensor of shape b t value value tensor of shape b s optional defaults to query returns mask a boolean tensor of shape 1 t s containing a lower triangular matrix of shape t s take the dot product between query and key to get the raw attention scores batch_dim query_heads target_seq_len source_seq_len this is actually dropping out entire tokens to attend to which might seem a bit unusual but is taken from the original transformer paper normalize the attention scores to probabilities scores b n t s the expand dim happens starting from the num_heads dimension batch_dims num_heads query_attention_dims key_attention_dims
from keras import constraints
from keras import initializers
from keras import ops
from keras import regularizers
from keras.api_export import keras_export
from keras.layers.activations.softmax import Softmax
from keras.layers.core.einsum_dense import EinsumDense
from keras.layers.layer import Layer
from keras.layers.regularization.dropout import Dropout


@keras_export("keras.layers.GroupQueryAttention")
class GroupedQueryAttention(Layer):
    

    def __init__(
        self,
        head_dim,
        num_query_heads,
        num_key_value_heads,
        dropout=0.0,
        use_bias=True,
        kernel_initializer="glorot_uniform",
        bias_initializer="zeros",
        kernel_regularizer=None,
        bias_regularizer=None,
        activity_regularizer=None,
        kernel_constraint=None,
        bias_constraint=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.supports_masking = True
        self.head_dim = head_dim
        self.num_query_heads = num_query_heads
        self.num_key_value_heads = num_key_value_heads
        if num_query_heads % num_key_value_heads != 0:
            raise ValueError(
                "`num_query_heads` must be divisible"
                " by `num_key_value_heads`."
            )
        self.num_repeats = num_query_heads // num_key_value_heads
        self.dropout = dropout
        self.use_bias = use_bias
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.bias_initializer = initializers.get(bias_initializer)
        self.kernel_regularizer = regularizers.get(kernel_regularizer)
        self.bias_regularizer = regularizers.get(bias_regularizer)
        self.activity_regularizer = regularizers.get(activity_regularizer)
        self.kernel_constraint = constraints.get(kernel_constraint)
        self.bias_constraint = constraints.get(bias_constraint)

    def build(
        self,
        query_shape,
        value_shape,
        key_shape=None,
    ):
        
        
        
        
        
        
        
        
        key_shape = value_shape if key_shape is None else key_shape
        self.feature_dim = query_shape[-1]
        self._query_dense = EinsumDense(
            "bqm,muh->bquh",
            output_shape=(None, self.num_query_heads, self.head_dim),
            bias_axes="uh" if self.use_bias else None,
            name="query",
            **self._get_common_kwargs_for_sublayer(),
        )
        self._query_dense.build(query_shape)

        self._key_dense = EinsumDense(
            "bkm,mvh->bkvh",
            output_shape=(None, self.num_key_value_heads, self.head_dim),
            bias_axes="vh" if self.use_bias else None,
            name="key",
            **self._get_common_kwargs_for_sublayer(),
        )
        self._key_dense.build(key_shape)

        self._value_dense = EinsumDense(
            "bkm,mvh->bkvh",
            output_shape=(None, self.num_key_value_heads, self.head_dim),
            bias_axes="vh" if self.use_bias else None,
            name="value",
            **self._get_common_kwargs_for_sublayer(),
        )
        self._value_dense.build(value_shape)

        self._softmax = Softmax(axis=-1, dtype=self.dtype_policy)
        self._dropout_layer = Dropout(
            rate=self.dropout, dtype=self.dtype_policy
        )

        self._dot_product_equation = "bquh,bkuh->buqk"
        self._combine_equation = "buqk,bkuh->bquh"

        self._output_dense = EinsumDense(
            "bquh,uhm->bqm",
            output_shape=(None, self.feature_dim),
            bias_axes="m" if self.use_bias else None,
            name="attention_output",
            **self._get_common_kwargs_for_sublayer(),
        )
        self._output_dense.build(
            (None, None, self.num_query_heads, self.head_dim)
        )
        self.built = True

    def _get_common_kwargs_for_sublayer(self):
        common_kwargs = dict(
            kernel_regularizer=self.kernel_regularizer,
            bias_regularizer=self.bias_regularizer,
            activity_regularizer=self.activity_regularizer,
            kernel_constraint=self.kernel_constraint,
            bias_constraint=self.bias_constraint,
            dtype=self.dtype_policy,
        )
        
        
        
        kernel_initializer = self.kernel_initializer.__class__.from_config(
            self.kernel_initializer.get_config()
        )
        bias_initializer = self.bias_initializer.__class__.from_config(
            self.bias_initializer.get_config()
        )
        common_kwargs["kernel_initializer"] = kernel_initializer
        common_kwargs["bias_initializer"] = bias_initializer
        return common_kwargs

    def call(
        self,
        query,
        value,
        key=None,
        query_mask=None,
        value_mask=None,
        key_mask=None,
        attention_mask=None,
        return_attention_scores=False,
        training=None,
        use_causal_mask=False,
    ):
        if key is None:
            key = value

        attention_mask = self._compute_attention_mask(
            query,
            value,
            query_mask=query_mask,
            value_mask=value_mask,
            key_mask=key_mask,
            attention_mask=attention_mask,
            use_causal_mask=use_causal_mask,
        )

        query = self._query_dense(query)
        key = self._key_dense(key)
        value = self._value_dense(value)

        key = ops.repeat(
            key, self.num_repeats, axis=2
        )  
        value = ops.repeat(
            value, self.num_repeats, axis=2
        )  

        output, scores = self._compute_attention(
            query,
            key,
            value,
            attention_mask=attention_mask,
            training=training,
        )

        output = self._output_dense(
            output
        )  

        if return_attention_scores:
            return output, scores
        return output

    def _compute_attention_mask(
        self,
        query,
        value,
        query_mask=None,
        value_mask=None,
        key_mask=None,
        attention_mask=None,
        use_causal_mask=False,
    ):
        
        auto_mask = None
        if query_mask is not None:
            query_mask = ops.cast(query_mask, "bool")  
            
            auto_mask = ops.expand_dims(query_mask, -1)  
        if value_mask is not None:
            value_mask = ops.cast(value_mask, "bool")  
            
            mask = ops.expand_dims(value_mask, -2)  
            auto_mask = mask if auto_mask is None else auto_mask & mask
        if key_mask is not None:
            key_mask = ops.cast(key_mask, "bool")  
            
            mask = ops.expand_dims(key_mask, -2)  
            auto_mask = mask if auto_mask is None else auto_mask & mask
        if use_causal_mask:
            
            mask = self._compute_causal_mask(query, value)
            auto_mask = mask if auto_mask is None else auto_mask & mask
        if auto_mask is not None:
            
            attention_mask = (
                auto_mask
                if attention_mask is None
                else ops.cast(attention_mask, bool) & auto_mask
            )
        return attention_mask

    def _compute_causal_mask(self, query, value=None):
        
        q_seq_length = ops.shape(query)[1]
        v_seq_length = q_seq_length if value is None else ops.shape(value)[1]
        ones_mask = ops.ones((1, q_seq_length, v_seq_length), dtype="int32")
        row_index = ops.cumsum(ones_mask, axis=-2)
        col_index = ops.cumsum(ones_mask, axis=-1)
        return ops.greater_equal(row_index, col_index)

    def _compute_attention(
        self, query, key, value, attention_mask=None, training=None
    ):
        query = ops.multiply(
            query,
            1.0 / ops.sqrt(ops.cast(self.head_dim, query.dtype)),
        )
        
        
        scores = ops.einsum(
            self._dot_product_equation, query, key
        )  
        scores = self._masked_softmax(scores, attention_mask=attention_mask)
        
        
        scores_dropout = self._dropout_layer(scores, training=training)
        output = ops.einsum(self._combine_equation, scores_dropout, value)
        return output, scores

    def _masked_softmax(self, scores, attention_mask=None):
        
        
        if attention_mask is not None:
            
            
            
            mask_expansion_axis = -1 * 2 - 1
            for _ in range(len(scores.shape) - len(attention_mask.shape)):
                attention_mask = ops.expand_dims(
                    attention_mask, axis=mask_expansion_axis
                )
        return self._softmax(scores, mask=attention_mask)

    def compute_output_shape(
        self,
        query_shape,
        value_shape,
        key_shape=None,
    ):
        if key_shape is None:
            key_shape = value_shape

        if query_shape[-1] != value_shape[-1]:
            raise ValueError(
                "The last dimension of `query_shape` and `value_shape` "
                f"must be equal, but are {query_shape[-1]}, {value_shape[-1]}. "
                "Received: query_shape={query_shape}, value_shape={value_shape}"
            )

        if value_shape[1:-1] != key_shape[1:-1]:
            raise ValueError(
                "All dimensions of `value` and `key`, except the last one, "
                f"must be equal. Received: value_shape={value_shape} and "
                f"key_shape={key_shape}"
            )

        return query_shape

    def get_config(self):
        config = {
            "head_dim": self.head_dim,
            "num_query_heads": self.num_query_heads,
            "num_key_value_heads": self.num_key_value_heads,
            "use_bias": self.use_bias,
            "dropout": self.dropout,
            "kernel_initializer": initializers.serialize(
                self.kernel_initializer
            ),
            "bias_initializer": initializers.serialize(self.bias_initializer),
            "kernel_regularizer": regularizers.serialize(
                self.kernel_regularizer
            ),
            "bias_regularizer": regularizers.serialize(self.bias_regularizer),
            "activity_regularizer": regularizers.serialize(
                self.activity_regularizer
            ),
            "kernel_constraint": constraints.serialize(self.kernel_constraint),
            "bias_constraint": constraints.serialize(self.bias_constraint),
        }
        base_config = super().get_config()
        return {**base_config, **config}