Julia and Zygote: Adjoint Error When Doing QR Factorization

Viewed 41

I have written a surrogate function called GEKPLS and I am trying to make the code work with an optimizer to find the optimal theta parameter values.

As a first step, I'm trying to make it work with Zygote and have the following code:

function min_rlfv(theta)
    g = GEKPLS(X, y, grads, n_comp, delta_x, xlimits, extra_points, theta)
    return -g.reduced_likelihood_function_value
end

Zygote.gradient(min_rlfv, [0.01, 0.1])

Running the above code results in the following error message:

ERROR: Need an adjoint for constructor LinearAlgebra.QRCompactWYQ{Float64, Matrix{Float64}}. Gradient is of type LinearAlgebra.Transpose{Float64, Matrix{Float64}}

The stacktrace leads to the following function. The highlighted line in my code editor is Q, G = qr(Ft) from the following code block:

function _reduced_likelihood_function(theta, kernel_type, d, nt, ij, y_norma, noise = 0.0)
    reduced_likelihood_function_value = -Inf
    nugget = 1000000.0 * eps() #a jitter for numerical stability; 
    if kernel_type == "squar_exp" 
        r = squar_exp(theta, d)
    end
    R = (I + zeros(nt, nt)) .* (1.0 + nugget + noise)
    for k in 1:size(ij)[1]
        R[ij[k, 1], ij[k, 2]] = r[k]
        R[ij[k, 2], ij[k, 1]] = r[k]
    end
    C = cholesky(R).L 
    F = ones(nt, 1) 
    Ft = C \ F
    Q, G = qr(Ft)
    Q = Array(Q)
    Yt = C \ y_norma
    beta = G \ [(transpose(Q) ⋅ Yt)]
    rho = Yt .- (Ft .* beta)
    gamma = transpose(C) \ rho
    sigma2 = sum((rho) .^ 2, dims = 1) / nt
    detR = prod(diag(C) .^ (2.0 / nt))
    reduced_likelihood_function_value = -nt * log10(sum(sigma2)) - nt * log10(detR)
    return beta, gamma, reduced_likelihood_function_value
end

Any pointers on how this can be fixed?

0 Answers
Related