function [X, W, M] = internal_attention(...
    Q,K,V,attnMask,numHeads,scale,keySize,valueSize,numQueries,numKeys,dropoutProb)
% This function is undocumented and will change in a future release.

% Applies attention.
%
% Inputs:
%   Q            - Query                 (N1) x (T1) x (N) array.
%   K            - Key                   (N1) x (T2) x (N) array.
%   V            - Value                 (N2) x (T2) x (N) array.
%   attnMask     - Attention mask        (T2) x (T1) x 1 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.
%   dropoutProb  - Dropout probability.
%
% Outputs:
%   X         - Output               (N2) x (T1) x (N) array.
%   W         - Attention weights    (T2) x (T1) x (numHeads) x (N) array.
%   M         - Dropout mask         (T2) x (T1) x (numHeads) x (N) array.

%   Copyright 2022 The MathWorks, Inc.

if ~isreal(Q) || ~isreal(K) || ~isreal(V)
    error(message('deep:dlarray:ComplexNotSupported'))
end

[X, W, M] = nnet.internal.cnngpu.attentionForward(Q, K, V, attnMask, ...
    numHeads, scale, keySize, valueSize, numQueries, numKeys, dropoutProb);
end