function layer = layerNormalizationLayer(NameValueArgs)
% layerNormalizationLayer   Layer normalization layer
%
%   layer = layerNormalizationLayer creates a layer normalization layer.
%   This type of layer normalizes the input data across channels and
%   applies an affine transformation with learnable parameters.
%
%   layer = layerNormalizationLayer('PARAM1', VAL1, ...)
%   specifies optional parameter name/value pairs for creating the layer:
%
%     'Name'                    - A name for the layer. Default is ''.
%
%     'Offset'                  - Initial value for learnable parameter
%                                 Offset (also called beta). This can be:
%                                    - []
%                                    - a 1-by-C array (for 1-D images)
%                                    - a 1-by-1-by-C array (for 2-D images)
%                                    - a 1-by-1-by-1-by-C array (for 3-D
%                                       images)
%                                    - a C-by-1 array (for sequence and
%                                       feature vectors)
%                                    where C is the number of channels.
%                                   The default is [].
%
%     'Scale'                   - Initial value for learnable parameter
%                                 Scale (also called gamma). This can be:
%                                    - []
%                                    - a 1-by-C array (for 1-D images)
%                                    - a 1-by-1-by-C array (for 2-D images)
%                                    - a 1-by-1-by-1-by-C array (for 3-D
%                                       images)
%                                    - a C-by-1 array (for sequence and
%                                       feature vectors)
%                                    where C is the number of channels.
%                                 The default is [].
%
%     'Epsilon'                 - Offset for the variance to avoid
%                                 divide-by-zero errors, specified as a
%                                 positive scalar.
%                                 The default is 1e-5.
%
%     'OperationDimension'      - Dimension to normalize over, specified as
%                                 one of these values: 
%                                 'auto'            - For feature,
%                                                     sequence, 1-D image, 
%                                                     or spatial-temporal
%                                                     data, normalize
%                                                     over the channel
%                                                     dimension. Otherwise,
%                                                     normalize over the
%                                                     spatial and channel
%                                                     dimensions.
%
%                                'channel-only'    - Normalize over the
%                                                    channel dimension.
%
%                                'spatial-channel' - Normalize over the
%                                                    spatial and channel
%                                                    dimensions.
%
%                                'batch-excluded'  - Normalize over all
%                                                    dimensions excluding
%                                                    the batch dimension.
%                                The default is 'auto'.
%
%     'OffsetInitializer'       - Function to initialize the Offset,
%                                 specified as 'narrow-normal', 'zeros',
%                                 'ones' or a function handle. The default
%                                 is 'zeros'.
%
%     'ScaleInitializer'        - Function to initialize the Scale,
%                                 specified as 'narrow-normal', 'ones' or
%                                 a function handle. The default is 'ones'.
%
%     'OffsetLearnRateFactor'   - Multiplier for the learning rate of
%                                 Offset. Default is 1.
%
%     'ScaleLearnRateFactor'    - Multiplier for the learning rate of
%                                 Scale. Default is 1.
%
%     'OffsetL2Factor'          - Multiplier for the L2 weight regulariser
%                                 for Offset. Default is 1.
%
%     'ScaleL2Factor'           - Multiplier for the L2 weight regulariser
%                                 for Scale. Default is 1.
%
%   Example:
%     % Create a layer normalization layer.
%
%     layer = layerNormalizationLayer;
%
%   See also nnet.cnn.layer.layerNormalizationLayer.
%
%   <a href="matlab:helpview('deeplearning','list_of_layers')">List of Deep Learning Layers</a>

%   Copyright 2020-2022 The MathWorks, Inc.

arguments
    NameValueArgs.Name                   {iAssertValidLayerName} = ""
    NameValueArgs.Epsilon                (1,1) = 1e-5
    NameValueArgs.Offset                 = []
    NameValueArgs.Scale                  = []
    NameValueArgs.OperationDimension     = 'auto'
    NameValueArgs.OffsetLearnRateFactor  (1,1) {iAssertValidFactor} = 1
    NameValueArgs.ScaleLearnRateFactor   (1,1) {iAssertValidFactor} = 1
    NameValueArgs.OffsetL2Factor         (1,1) {iAssertValidFactor} = 1
    NameValueArgs.ScaleL2Factor          (1,1) {iAssertValidFactor} = 1
    NameValueArgs.OffsetInitializer      = 'zeros'
    NameValueArgs.ScaleInitializer       = 'ones'
end

NameValueArgs.Epsilon = iValidateAndConvertEpsilon(NameValueArgs.Epsilon);
NameValueArgs.Name = convertStringsToChars(NameValueArgs.Name);
NameValueArgs.OperationDimension = iValidateOperationDimension(NameValueArgs.OperationDimension);

NameValueArgs = nnet.internal.cnn.layer.util.gatherParametersToCPU(NameValueArgs);

% Create an internal representation of the layer.
internalLayer = nnet.internal.cnn.layer.LayerNormalization(NameValueArgs.Name, [], NameValueArgs.Epsilon, ...
    NameValueArgs.OperationDimension);

internalLayer.Offset.L2Factor = NameValueArgs.OffsetL2Factor;
internalLayer.Offset.LearnRateFactor = NameValueArgs.OffsetLearnRateFactor;

internalLayer.Scale.L2Factor = NameValueArgs.ScaleL2Factor;
internalLayer.Scale.LearnRateFactor = NameValueArgs.ScaleLearnRateFactor;

% Use the internal layer to construct a user visible layer.
layer = nnet.cnn.layer.LayerNormalizationLayer(internalLayer);

% Set these properties (if supplied) using the visible layer so that the
% number of channels is correctly checked.
layer = iSetPropertyIfNonEmpty(layer, NameValueArgs, ...
    {'Offset', 'Scale' ...
    'OffsetInitializer', 'ScaleInitializer'});
end


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

function epsilon = iValidateAndConvertEpsilon(value)
epsilon = nnet.internal.cnn.layer.paramvalidation...
    .validateAndConvertEpsilon(value);
end

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

function layer = iSetPropertyIfNonEmpty(layer, args, propertyNames)
for i=1:numel(propertyNames)
    name = propertyNames{i};
    if ~isempty(args.(name))
        layer.(name) = args.(name);
    end
end
end

function dimString = iValidateOperationDimension(value)
dimString = validatestring(value, {'auto', 'channel-only', 'spatial-channel', 'batch-excluded'});
end