function layer = selfAttentionLayer( numHeads, numKeyChannels, nameValueArgs )
% selfAttentionLayer   Self attention layer
%
%   A self attention layer computes the queries, keys, and values by
%   performing linear transformations on the input. Then it uses them to
%   compute scaled dot-product attention across different heads, merges
%   the results from the different heads, and finally, performs a linear
%   transformation on the merged result.
%
%   layer = selfAttentionLayer(numHeads, numKeyChannels) creates a self
%   attention layer. numHeads specifies the number of attention heads.
%   numKeyChannels specifies the number of key and query channels.
%   numKeyChannels must be a positive integer and be divisible by numHeads.
%
%   layer = selfAttentionLayer(numHeads, numKeyChannels, Name=Value)
%   specifies additional options using one or more name-value arguments:
%
%     NumValueChannels         - Number of channels for the values,
%                                specified as a positive integer or "auto".
%                                If NumValueChannels is "auto", then the
%                                software uses numKeyChannels. The value
%                                must be divisible by numHeads. The default
%                                value is "auto".
%
%     OutputSize               - Number of channels for the layer output,
%                                specified as a positive integer or "auto".
%                                If OutputSize is "auto", then the software
%                                uses the number of channels in the layer
%                                input corresponding to the data. The
%                                default value is "auto".
%
%     HasPaddingMaskInput      - Flag indicating whether the layer has an
%                                extra input that represents the padding
%                                mask, specified as 1 (true) or 0 (false).
%                                If HasPaddingMaskInput is 1 (true), then
%                                the layer has an extra input with the name
%                                "mask". The padding mask input must be a
%                                logical array or numeric array consisting
%                                of 0 and 1. The software prevents and
%                                allows attention to elements in
%                                keys-values pairs when the corresponding
%                                element in the padding mask is 0 and 1
%                                respectively. The default value is 0
%                                (false).
%
%     AttentionMask            - Mask preventing attention to elements in
%                                certain positions in keys-values pairs,
%                                specified as one of the following:
%                                "none"   - Do not prevent attention to
%                                           elements with respect to their
%                                           positions. If
%                                           HasPaddingMaskInput is 1
%                                           (true), then the software
%                                           prevents attention to padding
%                                           elements only.
%                                "causal" - Prevent an element in position
%                                           M from attending to an element
%                                           in position N > M. Use this
%                                           option for auto-regressive
%                                           models.
%                                The default value is "none".
%
%     DropoutProbability       - Dropout probability for attention scores,
%                                specified as a nonnegative scalar less
%                                than 1. The default value is 0.
%
%     HasScoresOutput          - Flag indicating whether the layer has an
%                                extra output that represents attention
%                                scores, specified as 1 (true) or 0
%                                (false). If HasScoresOutput is 1 (true),
%                                then the layer has an output with the name
%                                "scores". The default value is 0 (false).
%
%     QueryWeights             - Query weights, specified as a 
%                                numKeyChannels-by-numInputChannels matrix
%                                or [], where numInputChannels is the
%                                number of channels in the layer input that
%                                corresponds to the data. The default value
%                                is [].
%
%     KeyWeights               - Key weights, specified as a 
%                                numKeyChannels-by-numInputChannels matrix
%                                or [], where numInputChannels is the
%                                number of channels in the layer input that
%                                corresponds to the data. The default value
%                                is [].
%
%     ValueWeights             - Value weights, specified as a 
%                                NumValueChannels-by-numInputChannels
%                                matrix or [], where numInputChannels is
%                                the number of channels in the layer input
%                                that corresponds to the data. The default
%                                value is [].
%
%     OutputWeights            - Layer output weights, specified as a 
%                                OutputSize-by-NumValueChannels matrix or
%                                []. The default value is [].
%
%     QueryBias                - Query biases, specified as a
%                                numKeyChannels-by-1 vector or []. The
%                                default value is [].
%
%     KeyBias                  - Key biases, specified as a
%                                numKeyChannels-by-1 vector or []. The
%                                default value is [].
%
%     ValueBias                - Value biases, specified as an
%                                NumValueChannels-by-1 vector or []. The
%                                default value is [].
%
%     OutputBias               - Output biases, specified as an
%                                OutputSize-by-1 vector or []. The default
%                                value is [].
%
%     WeightLearnRateFactor    - Multiplier for the learning rate for
%                                QueryWeights, KeyWeights, ValueWeights,
%                                and OutputWeights. The default value is 1.
%
%     BiasLearnRateFactor      - Multiplier for the learning rate for
%                                QueryBias, KeyBias, ValueBias, and
%                                OutputBias. The default value is 1.
%
%     WeightL2Factor           - Multiplier for the L2 weight regularizer
%                                for QueryWeights, KeyWeights,
%                                ValueWeights, and OutputWeights. The
%                                default value is 1.
%
%     BiasL2Factor             - Multiplier for the L2 weight regularizer
%                                for QueryBias, KeyBias, ValueBias, and
%                                OutputBias. The default value is 0.
%
%     WeightsInitializer       - Function to initialize QueryWeights,
%                                KeyWeights, ValueWeights, and
%                                OutputWeights, specified as "glorot",
%                                "he", "orthogonal", "narrow-normal",
%                                "zeros", "ones", or a function handle. The
%                                default value is "glorot".
%
%     BiasInitializer          - Function to initialize QueryBias, KeyBias,
%                                ValueBias, and OutputBias, specified as
%                                "narrow-normal", "zeros", "ones", or a
%                                function handle. The default value is
%                                "zeros".
%
%     Name                     - Name for the layer. The default value is
%                               "".
%
%   Example:
%       % Create a self attention layer with 8 heads and 256 key and query
%       % channels.
%
%       layer = selfAttentionLayer(8, 256);
%
%   See also nnet.cnn.layer.SelfAttentionLayer, lstmLayer,
%   convolution2dLayer.
%
%   <a href="matlab:helpview('deeplearning','list_of_layers')">List of Deep Learning Layers</a>

%   Copyright 2022-2023 The MathWorks, Inc.

arguments
    numHeads                             (1,1) {mustBeNumeric, mustBeInteger, mustBeNonempty, mustBePositive}
    numKeyChannels                       (1,1) {mustBeNumeric, mustBeInteger, mustBeNonempty, mustBePositive}
    nameValueArgs.Name                   {iAssertValidLayerName} = ""
    nameValueArgs.NumValueChannels       {iAssertValidNumChannels(nameValueArgs.NumValueChannels, 'NumValueChannels')} = 'auto'
    nameValueArgs.OutputSize             {iAssertValidNumChannels(nameValueArgs.OutputSize, 'OutputSize')} = 'auto'
    nameValueArgs.HasPaddingMaskInput    {iAssertBinary} = false
    nameValueArgs.HasScoresOutput        {iAssertBinary} = false
    nameValueArgs.AttentionMask          {mustBeTextScalar} = 'none'
    nameValueArgs.DropoutProbability     {iAssertValidDropoutProbability} = 0
    nameValueArgs.WeightLearnRateFactor {iAssertValidFactor} = 1
    nameValueArgs.BiasLearnRateFactor    {iAssertValidFactor} = 1
    nameValueArgs.WeightsInitializer     = 'glorot'
    nameValueArgs.BiasInitializer        = 'zeros'
    nameValueArgs.WeightL2Factor        {iAssertValidFactor} = 1
    nameValueArgs.BiasL2Factor           {iAssertValidFactor} = 0
    nameValueArgs.QueryWeights           = []
    nameValueArgs.KeyWeights             = []
    nameValueArgs.ValueWeights           = []
    nameValueArgs.OutputWeights          = []
    nameValueArgs.QueryBias              = []
    nameValueArgs.KeyBias                = []
    nameValueArgs.ValueBias              = []
    nameValueArgs.OutputBias             = []
end

% Store all input arguments in one struct
args = nameValueArgs;
args.NumHeads = numHeads;
args.NumKeyChannels = numKeyChannels;

% Gather arguments to CPU and convert them to canonical form
args = nnet.internal.cnn.layer.util.gatherParametersToCPU(args);
args = iConvertToCanonicalForm(args);
iAssertNumHeadsMatchesHiddenSizes(args.NumHeads, args.NumKeyChannels, args.NumValueChannels);

% Create an internal representation of a self attention layer
internalLayer = nnet.internal.cnn.layer.SelfAttention( ...
    args.Name, ...
    args.NumHeads, ...
    args.NumKeyChannels, ...
    args.NumValueChannels, ...
    args.OutputSize, ...
    args.InputSize, ...
    args.HasPaddingMaskInput, ...
    args.HasScoresOutput, ...
    args.AttentionMask, ...
    args.DropoutProbability);

internalLayer = internalLayer.setL2Factor(args.WeightL2Factor, args.BiasL2Factor);
internalLayer = internalLayer.setLearnRateFactor(args.WeightLearnRateFactor, args.BiasLearnRateFactor);

% Pass the internal layer to a function to construct a user visible
% multihead attention layer.
layer = nnet.cnn.layer.SelfAttentionLayer(internalLayer);
layer.WeightsInitializer = args.WeightsInitializer;
layer.BiasInitializer = args.BiasInitializer;
layer.QueryWeights = args.QueryWeights;
layer.KeyWeights = args.KeyWeights;
layer.ValueWeights = args.ValueWeights;
layer.OutputWeights = args.OutputWeights;
layer.QueryBias = args.QueryBias;
layer.KeyBias = args.KeyBias;
layer.ValueBias = args.ValueBias;
layer.OutputBias = args.OutputBias;
end

function iAssertNumHeadsMatchesHiddenSizes(numHeads, numKeyChannels, numValueChannels)
iAssertHiddenSize(numHeads, numKeyChannels, ...
    'nnet_cnn:internal:cnn:layer:SelfAttention:NumKeyChannelsNumHeadsMismatch');
if ~isempty(numValueChannels)
    iAssertHiddenSize(numHeads, numValueChannels, ...
        'nnet_cnn:internal:cnn:layer:SelfAttention:NumValueChannelsNumHeadsMismatch');
end
end

function iAssertHiddenSize(numHeads, hiddenSize, errorId)
if mod(hiddenSize, numHeads) ~= 0
    error( message(errorId, numHeads, hiddenSize) );
end
end

function iAssertValidNumChannels(value, type)
nnet.internal.cnn.layer.paramvalidation.validateNumChannels(value, type);
end

function iAssertValidFactor(value)
nnet.internal.cnn.layer.paramvalidation.validateLearnFactor(value);
end

function iAssertValidLayerName(name)
nnet.internal.cnn.layer.paramvalidation.validateLayerName(name)
end

function iAssertBinary(value)
nnet.internal.cnn.options.OptionsValidator.assertBinary(value);
end

function validString = iAssertValidAttentionMask(value)
validString = validatestring(value, {'none', 'causal'});
end

function iAssertValidDropoutProbability(value)
validateattributes(value, {'numeric'}, {'scalar','real','finite','nonsparse','nonnegative','<', 1});
end

function inputArguments = iConvertToCanonicalForm(params)
% Make sure integral values are converted to double and strings to char vectors
inputArguments = struct;
inputArguments.InputSize = [];
inputArguments.NumHeads = iConvertToDouble(params.NumHeads);
inputArguments.NumKeyChannels = iConvertToDouble( iConvertToEmptyIfAuto(params.NumKeyChannels) );
inputArguments.NumValueChannels = iConvertToDouble( iConvertToEmptyIfAuto(params.NumValueChannels) );
inputArguments.OutputSize = iConvertToDouble( iConvertToEmptyIfAuto(params.OutputSize) );
inputArguments.HasPaddingMaskInput = logical( params.HasPaddingMaskInput );
inputArguments.HasScoresOutput = logical( params.HasScoresOutput );
inputArguments.AttentionMask = iAssertValidAttentionMask(params.AttentionMask);
inputArguments.DropoutProbability = iConvertToDouble(params.DropoutProbability);
inputArguments.WeightLearnRateFactor = params.WeightLearnRateFactor;
inputArguments.BiasLearnRateFactor = params.BiasLearnRateFactor;
inputArguments.WeightL2Factor = params.WeightL2Factor;
inputArguments.BiasL2Factor = params.BiasL2Factor;
inputArguments.WeightsInitializer = params.WeightsInitializer;
inputArguments.BiasInitializer = params.BiasInitializer;
inputArguments.Name = char( params.Name );
inputArguments.QueryWeights = params.QueryWeights;
inputArguments.KeyWeights = params.KeyWeights;
inputArguments.ValueWeights = params.ValueWeights;
inputArguments.OutputWeights = params.OutputWeights;
inputArguments.QueryBias = params.QueryBias;
inputArguments.KeyBias = params.KeyBias;
inputArguments.ValueBias = params.ValueBias;
inputArguments.OutputBias = params.OutputBias;
end

function x = iConvertToEmptyIfAuto(x)
if strcmp(x, "auto")
  x = [];
end
end

function value = iConvertToDouble(value)
value = nnet.internal.cnn.layer.util.convertToDouble(value);
end
