How to get diagonal element in a custom layer in NN tool box?

function Uout = forward(layer, Uin)
    W = Uin(:,:,1) * layer.C;
    W = Uin(:,:,1) + repmat(layer.alpha .* diag(W*Uin(:,:,1)'),1,layer.usize(2)) ...
        .* Uin(:,:,1) + layer.B1 * Uin(:,:,1) - repmat(layer.alpha,1,layer.usize(2)).* Uin(:,:,1);
    Wnorm = norms(W,2,2);
    Uout = W./repmat(Wnorm,1,layer.usize(2));
end

I define the forward function like this, but it is said that dlarray has no compatible function diag, then how I solve this problem? I didnt find answer in the official documents.

Answer 1

[Aiswarya]

Hi,

I understand that you are trying to extract the diagonal elements of dlarray using diag function. The diag function can't be used with the dlarray datatype and there are no alternative functions provided by dlarray as well. However, there is a work around to get the diagonal elements by using the following script

diagonal_elements = W(logical(eye(size(W))));

The eye function (https://www.mathworks.com/help/matlab/ref/eye.html) creates an identity matrix of input size. On passing the size of weights matrix, the eye function will create identity matrix of that size which can be used as a logical index to obtain only the diagonal elements of weight matrix W as a dlarray vector.