multiheadattention layer this is an implementation of multiheaded attention as described in the paper attention is all you need vaswani et al 2017https arxiv orgabs1706 03762 if query key value are the same then this is selfattention each timestep in query attends to the corresponding sequence in key and returns a fixedwidth vector this layer first projects query key and value these are effectively a list of tensors of length numattentionheads where the corresponding shapes are batchsize query dimensions keydim batchsize keyvalue dimensions keydim batchsize keyvalue dimensions valuedim then the query and key tensors are dotproducted and scaled these are softmaxed to obtain attention probabilities the value tensors are then interpolated by these probabilities then concatenated back to a single tensor finally the result tensor with the last dimension as valuedim can take a linear projection and return args numheads number of attention heads keydim size of each attention head for query and key valuedim size of each attention head for value dropout dropout probability usebias boolean whether the dense layers use bias vectorsmatrices outputshape the expected shape of an output tensor besides the batch and sequence dims if not specified projects back to the query feature dim the query input s last dimension attentionaxes axes over which the attention is applied none means attention over all axes but batch heads and features 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 b t dim where b is the batch size t is the target sequence length and dim is the feature dimension value value tensor of shape b s dim where b is the batch size s is the source sequence length and dim is the feature dimension key optional key tensor of shape b s dim if not given will use value for both key and value which is the most common case attentionmask a boolean mask of shape b t s that prevents attention to certain positions the boolean mask specifies which query elements can attend to which key elements 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 the result of the computation of shape b t e where t is for target sequence shapes and e is the query input last dimension if outputshape is none otherwise the multihead outputs are projected to the shape specified by outputshape attentionscores optional multihead attention coefficients over attention axes cache 1 0 math sqrtself keydim builds layers and variables args queryshape shape of the query tensor valueshape shape of the value tensor key optional shape of the key tensor builds the attention computations for multihead dot product attention these computations could be wrapped into the keras attention layer once it supports multihead einsum computations 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 builds the output projection matrix args freedims number of free dimensions for einsum equation building commonkwargs common keyword arguments for einsum layer name name for the projection layer returns projection layer builds multihead dotproduct attention computations this function builds attributes necessary for computeattention to customize attention computation to replace the default dotproduct attention args rank the rank of query key value tensors normalize the attention scores to probabilities attentionscores b n t s the expand dim happens starting from the numheads dimension batchdims numheads queryattentiondims keyattentiondims applies dotproduct attention with query key value tensors this function defines the computation inside call with projected multihead q k v inputs users can override this function for customized attention implementation args query projected query tensor of shape b t n keydim key projected key tensor of shape b s n keydim value projected value tensor of shape b s n valuedim attentionmask a boolean mask of shape b t s that prevents attention to certain positions it is generally not needed if the query and value andor key are masked training python boolean indicating whether the layer should behave in training mode adding dropout or in inference mode doing nothing returns attentionoutput multiheaded outputs of attention computation attentionscores multiheaded attention weights note applying scalar multiply at the smaller end of einsum improves xla performance but may introduce slight numeric differences in the transformer attention head 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 contextlayer b t n h n numattentionheads h sizeperhead query b t n h key b s n h value b s n h 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 coverts an index to a einsum variable name we simply map indices to lowercase characters e g 0 a 1 b builds einsum equations for the attention computation query key value inputs after projection are expected to have the shape as bs nonattention dims attention dims numheads channels bs and nonattention dims are treated as batch dims the attention operations can be generalized 1 querykey dot product batch dims query attention dims numheads channels batch dims key attention dims numheads channels batch dims numheads query attention dims key attention dims 2 combination batch dims numheads query attention dims key attention dims batch dims value attention dims numheads channels batch dims query attention dims numheads channels args rank rank of query key value tensors attnaxes listtuple of axes 1 rank that attention will be applied to returns einsum equations batchdims includes the head dim builds an einsum equation for projections inside multihead attention inputstr kernelstr outputstr biasaxes letteroffset 0 for i in rangefreedims char indextoeinsumvariablei letteroffset inputstr char outputstr char letteroffset freedims for i in rangebounddims char indextoeinsumvariablei letteroffset inputstr char kernelstr char letteroffset bounddims for i in rangeoutputdims char indextoeinsumvariablei letteroffset kernelstr char outputstr char biasaxes char equation finputstr kernelstroutputstr return equation biasaxes lenoutputstr def getoutputshapeoutputrank knownlastdims return none outputrank lenknownlastdims listknownlastdims multiheadattention layer this is an implementation of multi headed attention as described in the paper attention is all you need vaswani et al 2017 https arxiv org abs 1706 03762 if query key value are the same then this is self attention each timestep in query attends to the corresponding sequence in key and returns a fixed width vector this layer first projects query key and value these are effectively a list of tensors of length num_attention_heads where the corresponding shapes are batch_size query dimensions key_dim batch_size key value dimensions key_dim batch_size key value dimensions value_dim then the query and key tensors are dot producted and scaled these are softmaxed to obtain attention probabilities the value tensors are then interpolated by these probabilities then concatenated back to a single tensor finally the result tensor with the last dimension as value_dim can take a linear projection and return args num_heads number of attention heads key_dim size of each attention head for query and key value_dim size of each attention head for value dropout dropout probability use_bias boolean whether the dense layers use bias vectors matrices output_shape the expected shape of an output tensor besides the batch and sequence dims if not specified projects back to the query feature dim the query input s last dimension attention_axes axes over which the attention is applied none means attention over all axes but batch heads and features 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 b t dim where b is the batch size t is the target sequence length and dim is the feature dimension value value tensor of shape b s dim where b is the batch size s is the source sequence length and dim is the feature dimension key optional key tensor of shape b s dim if not given will use value for both key and value which is the most common case attention_mask a boolean mask of shape b t s that prevents attention to certain positions the boolean mask specifies which query elements can attend to which key elements 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 the result of the computation of shape b t e where t is for target sequence shapes and e is the query input last dimension if output_shape is none otherwise the multi head outputs are projected to the shape specified by output_shape attention_scores optional multi head attention coefficients over attention axes cache 1 0 math sqrt self _key_dim builds layers and variables args query_shape shape of the query tensor value_shape shape of the value tensor key optional shape of the key tensor builds the attention computations for multi head dot product attention these computations could be wrapped into the keras attention layer once it supports multi head einsum computations 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 builds the output projection matrix args free_dims number of free dimensions for einsum equation building common_kwargs common keyword arguments for einsum layer name name for the projection layer returns projection layer builds multi head dot product attention computations this function builds attributes necessary for _compute_attention to customize attention computation to replace the default dot product attention args rank the rank of query key value tensors normalize the attention scores to probabilities attention_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 applies dot product attention with query key value tensors this function defines the computation inside call with projected multi head q k v inputs users can override this function for customized attention implementation args query projected query tensor of shape b t n key_dim key projected key tensor of shape b s n key_dim value projected value tensor of shape b s n value_dim attention_mask a boolean mask of shape b t s that prevents attention to certain positions it is generally not needed if the query and value and or key are masked training python boolean indicating whether the layer should behave in training mode adding dropout or in inference mode doing nothing returns attention_output multi headed outputs of attention computation attention_scores multi headed attention weights note applying scalar multiply at the smaller end of einsum improves xla performance but may introduce slight numeric differences in the transformer attention head 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 context_layer b t n h n num_attention_heads h size_per_head query b t n h key b s n h value b s n h 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 coverts an index to a einsum variable name we simply map indices to lowercase characters e g 0 a 1 b builds einsum equations for the attention computation query key value inputs after projection are expected to have the shape as bs non attention dims attention dims num_heads channels bs and non attention dims are treated as batch dims the attention operations can be generalized 1 query key dot product batch dims query attention dims num_heads channels batch dims key attention dims num_heads channels batch dims num_heads query attention dims key attention dims 2 combination batch dims num_heads query attention dims key attention dims batch dims value attention dims num_heads channels batch dims query attention dims num_heads channels args rank rank of query key value tensors attn_axes list tuple of axes 1 rank that attention will be applied to returns einsum equations batch_dims includes the head dim builds an einsum equation for projections inside multi head attention
import collections
import math
import string

import numpy as np

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.MultiHeadAttention")
class MultiHeadAttention(Layer):
    

    def __init__(
        self,
        num_heads,
        key_dim,
        value_dim=None,
        dropout=0.0,
        use_bias=True,
        output_shape=None,
        attention_axes=None,
        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._num_heads = num_heads
        self._key_dim = key_dim
        
        self._inverse_sqrt_key_dim = None
        self._value_dim = value_dim if value_dim else key_dim
        self._dropout = dropout
        self._use_bias = use_bias
        self._output_shape = output_shape
        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)
        if isinstance(attention_axes, int):
            attention_axes = (attention_axes,)
        elif attention_axes and not isinstance(attention_axes, (list, tuple)):
            raise ValueError(
                "`attention_axes` must be an int, list, or tuple."
                f"Received: attention_axes={attention_axes}"
            )
        self._attention_axes = attention_axes

    def get_config(self):
        base_config = super().get_config()
        config = {
            "num_heads": self._num_heads,
            "key_dim": self._key_dim,
            "value_dim": self._value_dim,
            "dropout": self._dropout,
            "use_bias": self._use_bias,
            "output_shape": self._output_shape,
            "attention_axes": self._attention_axes,
            "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),
        }
        return {**base_config, **config}

    def build(
        self,
        query_shape,
        value_shape,
        key_shape=None,
    ):
        
        key_shape = value_shape if key_shape is None else key_shape
        query_rank = len(query_shape)
        value_rank = len(value_shape)
        key_rank = len(key_shape)
        einsum_equation, bias_axes, output_rank = _build_proj_equation(
            query_rank - 1, bound_dims=1, output_dims=2
        )
        self._query_dense = EinsumDense(
            einsum_equation,
            output_shape=_get_output_shape(
                output_rank - 1, [self._num_heads, self._key_dim]
            ),
            bias_axes=bias_axes if self._use_bias else None,
            name="query",
            **self._get_common_kwargs_for_sublayer(),
        )
        self._query_dense.build(query_shape)
        einsum_equation, bias_axes, output_rank = _build_proj_equation(
            key_rank - 1, bound_dims=1, output_dims=2
        )
        self._key_dense = EinsumDense(
            einsum_equation,
            output_shape=_get_output_shape(
                output_rank - 1, [self._num_heads, self._key_dim]
            ),
            bias_axes=bias_axes if self._use_bias else None,
            name="key",
            **self._get_common_kwargs_for_sublayer(),
        )
        self._key_dense.build(key_shape)
        einsum_equation, bias_axes, output_rank = _build_proj_equation(
            value_rank - 1, bound_dims=1, output_dims=2
        )
        self._value_dense = EinsumDense(
            einsum_equation,
            output_shape=_get_output_shape(
                output_rank - 1, [self._num_heads, self._value_dim]
            ),
            bias_axes=bias_axes if self._use_bias else None,
            name="value",
            **self._get_common_kwargs_for_sublayer(),
        )
        self._value_dense.build(value_shape)

        
        
        
        self._build_attention(output_rank)
        self._output_dense = self._make_output_dense(
            query_shape,
            self._get_common_kwargs_for_sublayer(),
            "attention_output",
        )
        output_dense_input_shape = list(
            self._query_dense.compute_output_shape(query_shape)
        )
        output_dense_input_shape[-1] = self._value_dim
        self._output_dense.build(tuple(output_dense_input_shape))
        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 _make_output_dense(self, query_shape, common_kwargs, name=None):
        
        query_rank = len(query_shape)
        if self._output_shape:
            if not isinstance(self._output_shape, collections.abc.Sized):
                output_shape = [self._output_shape]
            else:
                output_shape = self._output_shape
        else:
            output_shape = [query_shape[-1]]
        einsum_equation, bias_axes, output_rank = _build_proj_equation(
            query_rank - 1, bound_dims=2, output_dims=len(output_shape)
        )
        return EinsumDense(
            einsum_equation,
            output_shape=_get_output_shape(output_rank - 1, output_shape),
            bias_axes=bias_axes if self._use_bias else None,
            name=name,
            **common_kwargs,
        )

    def _build_attention(self, rank):
        
        if self._attention_axes is None:
            self._attention_axes = tuple(range(1, rank - 2))
        else:
            self._attention_axes = tuple(self._attention_axes)
        (
            self._dot_product_equation,
            self._combine_equation,
            attn_scores_rank,
        ) = _build_attention_equation(rank, attn_axes=self._attention_axes)
        norm_axes = tuple(
            range(
                attn_scores_rank - len(self._attention_axes), attn_scores_rank
            )
        )
        self._softmax = Softmax(axis=norm_axes, dtype=self.dtype_policy)
        self._dropout_layer = Dropout(
            rate=self._dropout, dtype=self.dtype_policy
        )
        self._inverse_sqrt_key_dim = 1.0 / math.sqrt(float(self._key_dim))

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

    def _compute_attention(
        self, query, key, value, attention_mask=None, training=None
    ):
        
        
        
        
        query = ops.multiply(
            query, ops.cast(self._inverse_sqrt_key_dim, query.dtype)
        )

        
        
        attention_scores = ops.einsum(self._dot_product_equation, key, query)

        attention_scores = self._masked_softmax(
            attention_scores, attention_mask
        )

        
        
        attention_scores_dropout = self._dropout_layer(
            attention_scores, training=training
        )

        
        attention_output = ops.einsum(
            self._combine_equation, attention_scores_dropout, value
        )
        return attention_output, attention_scores

    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)

        attention_output, attention_scores = self._compute_attention(
            query, key, value, attention_mask, training
        )
        attention_output = self._output_dense(attention_output)

        if return_attention_scores:
            return attention_output, attention_scores
        return attention_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_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}"
            )

        if self._output_shape:
            return query_shape[:-1] + self._output_shape

        return query_shape


def _index_to_einsum_variable(i):
    
    return string.ascii_lowercase[i]


def _build_attention_equation(rank, attn_axes):
    
    target_notation = ""
    for i in range(rank):
        target_notation += _index_to_einsum_variable(i)
    
    batch_dims = tuple(np.delete(range(rank), attn_axes + (rank - 1,)))
    letter_offset = rank
    source_notation = ""
    for i in range(rank):
        if i in batch_dims or i == rank - 1:
            source_notation += target_notation[i]
        else:
            source_notation += _index_to_einsum_variable(letter_offset)
            letter_offset += 1

    product_notation = "".join(
        [target_notation[i] for i in batch_dims]
        + [target_notation[i] for i in attn_axes]
        + [source_notation[i] for i in attn_axes]
    )
    dot_product_equation = "%s,%s->%s" % (
        source_notation,
        target_notation,
        product_notation,
    )
    attn_scores_rank = len(product_notation)
    combine_equation = "%s,%s->%s" % (
        product_notation,
        source_notation,
        target_notation,
    )
    return dot_product_equation, combine_equation, attn_scores_rank


def _build_proj_equation(free_dims, bound_dims, output_dims):
    
    input_str = ""
    kernel_str = ""
    output_str = ""
    bias_axes = ""
    letter_offset = 0
    for i in range(free_dims):
        char = _index_to_einsum_variable(i + letter_offset)
        input_str += char
        output_str += char

    letter_offset += free_dims
    for i in range(bound_dims):
        char = _index_to_einsum_variable(i + letter_offset)
        input_str += char
        kernel_str += char

    letter_offset += bound_dims
    for i in range(output_dims):
        char = _index_to_einsum_variable(i + letter_offset)
        kernel_str += char
        output_str += char
        bias_axes += char
    equation = f"{input_str},{kernel_str}->{output_str}"

    return equation, bias_axes, len(output_str)


def _get_output_shape(output_rank, known_last_dims):
    return [None] * (output_rank - len(known_last_dims)) + list(known_last_dims)