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.

%   Copyright 2019-2021 The MathWorks, Inc.

% gpuArray method for internal_batchnormBackward. The help for this method
% is similar to the CPU method internal_batchnormBackward

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

% Ensure all inputs have the same data type
[dZ, X, scale, batchMean, batchInvSqrtVar] = ...
    nnet.internal.cnngpu.castToSuperiorFloat(dZ, X, scale, batchMean, batchInvSqrtVar);

if (computationMode == "inference")
    % NOTE: We are using a cnnhost method (implemented in MATLAB) for
    % gpuArray because the GPU builtin computes gradients in training mode
    % only.
    [dZ, dOffset, dScale] = nnet.internal.cnnhost.batchNormalizationBackwardInference(...
        dZ, X, scale, epsilon, batchMean, batchInvSqrtVar, channelDim);
else
    % Find the reshape needed for the builtin, and reshape if needed
    [newSize, originalSize, channelDim, needsShapeChange, paramSize] = ...
        deep.internal.dlarray.makeValidSizeChannelWiseOperation(X, channelDim, size(X,channelDim));
    if needsShapeChange
        X = reshape(X, newSize);
        dZ = reshape(dZ, newSize);
        
        % Reshape the parameters to match the new input shape
        originalSizeParams = size(scale);
        scale = reshape(scale, paramSize);
        batchMean = reshape(batchMean, paramSize);
        batchInvSqrtVar = reshape(batchInvSqrtVar, paramSize);
    end
    
    % Call the CUDNN builtin
    [dZ, dOffset, dScale] = nnet.internal.cnngpu.batchNormalizationBackward(...
        dZ, X, scale, epsilon, batchMean, batchInvSqrtVar, channelDim);
    
    % Restore the data shape
    if needsShapeChange
        dZ = reshape(dZ, originalSize);
        dOffset = reshape(dOffset, originalSizeParams);
        dScale = reshape(dScale, originalSizeParams);
    end
end
end
