function Z = make_diag_real(A)
    %MAKE_DIAG_REAL Summary of this function goes here
    %   Detailed explanation goes here
    sz = size(A);
    
    assert(sz(1) == sz(2));

    if length(sz) == 2
        Z = A;
        Z(1:sz(1)+1:sz(1)*sz(2)) = real(A(1:sz(1)+1:sz(1)*sz(2)));
    elseif length(sz) == 3
        Z = reshape(A, sz(1)*sz(2), []);
        Z(1:sz(1)+1:sz(1)*sz(2),:) = real(Z(1:sz(1)+1:sz(1)*sz(2),:));
        Z = reshape(Z, sz);
    else
        error("Shape not supported")
    end
end

