function [dQ,dK,dV] = internal_attentionBackward(...
    dZ,Q,K,V,W,dropoutMask,numHeads,scale,keySize,valueSize,numQueries,numKeys)
% This function is undocumented and will change in a future release.

% Backward method for attention.
%
% Inputs:
%  dZ           - Derivative w.r.t the function   (N2) x (T1) x (N) array.
%  Q            - Query                           (N1) x (T1) x (N) array.
%  K            - Key                             (N1) x (T2) x (N) array.
%  V            - Value                           (N2) x (T2) x (N) array.
%  W            - Attention weights               (T2) x (T1) x (numHeads) x (N) array.
%  dropoutMask  - Dropout mask                    (T2) x (T1) x (numHeads) x (N) array.
%  numHeads     - Number of attention heads. Must divide N1 and N2.
%  scale        - Scale factor.
%  keySize      - Key size               = N1/numHeads.
%  valueSize    - Value size             = N2/numheads.
%  numQueries   - Number of queries = T1.
%  numKeys      - Number of keys = T2.
%
% Outputs:
%  dQ           - Derivative w.r.t query  (N1) x (T1) x (N) array.
%  dK           - Derivative w.r.t key    (N1) x (T2) x (N) array.
%  dV           - Derivative w.r.t value  (N2) x (T2) x (N) array.

%   Copyright 2022 The MathWorks, Inc.

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

% Split attention heads
Q = deep.internal.dlarray.splitAttentionHeads(Q, numHeads, keySize, numQueries);
K = deep.internal.dlarray.splitAttentionHeads(K, numHeads, keySize, numKeys);
V = deep.internal.dlarray.splitAttentionHeads(V, numHeads, valueSize, numKeys);

[dQ, dK, dV] = nnet.internal.cnngpu.attentionBackward(dZ, Q, K, V, W, ...
    dropoutMask, numHeads, scale, keySize, valueSize, numQueries, numKeys);

end