function [dZ, dOffset, dScale] = internal_batchnormBackward(dZ, X, scale, epsilon, batchMean, batchInvSqrtVar, channelDim, computationMode)
% This function is undocumented and will change in a future release.

%   [dX, dOffset, dScale] = internal_batchnormBackward(dZ, X, gamma, epsilon, batchMean, batchInvSqrtVar, channelDim, computationMode)
%   Compute adjoints for batch normalization.
%
%   Inputs:
%     dZ              Adjoint of forward result. Shape is only non-trivial in channel dimension.
%                     numeric array
%     X               Batch data.
%                     numeric array
%     scale           Scale. Also called gamma. Shape is only non-trivial in channel dimension.
%                     numeric array
%     epsilon         Constant to avoid division by zero.
%                     numeric scalar
%     batchMean       Batch mean. Shape is only non-trivial in channel dimension.
%                     numeric array
%     batchInvSqrtVar Inverse-sqrt-variance. Shape is only non-trivial in channel dimension.
%                     numeric array
%     channelDim      Channel dimension.
%                     double
%     computationMode Computation mode - either "training" or "inference".
%                     NOTE: In inference mode, batchMean is the trainedMean
%                     and batchInvSqrtVar is the trainedInvSqrtVar which is
%                     defined as 1./sqrt(trainedVariance + epsilon).
%
%   Outputs:
%     dX              Adjoint of X. Shape is same as X.
%                     numeric array
%     dOffset           Adjoint of offset. Shape is same as scale.
%                     numeric array
%     dScale          Adjoint of scale. Shape is same as scale.
%                     numeric array

%   Copyright 2019-2021 The MathWorks, Inc.

if ~isreal(dZ)
    error(message('deep:dlarray:ComplexGradientOutput', 'batchnorm'));
end

if (computationMode == "inference")
    [dZ, dOffset, dScale] = nnet.internal.cnnhost.batchNormalizationBackwardInference(...
        dZ, X, scale, epsilon, batchMean, batchInvSqrtVar, channelDim);
else
    [dZ, dOffset, dScale] = nnet.internal.cnnhost.batchNormalizationBackward(...
        dZ, X, scale, epsilon, batchMean, batchInvSqrtVar, channelDim);
end
end
