I’m trying to write a custom adjoint for a function given below:
function factor_divide(poly::Vector{ComplexF64}, α::ComplexF64, p::Int64)
quotient_deg::Int64 = size(poly,1) - p - 1
quotient::Vector{ComplexF64} = zeros(ComplexF64, quotient_deg + 1)
g::Vector{ComplexF64} = map(i->(-α)^(p-i+1)*binomial(p,i-1),1:p+1)
quotient[1] = poly[1]/g[1]
for i in 2:size(quotient, 1)
quotient[i] = (poly[i]-sum(j->quotient[j]*g[i-j+1],max(1,i-p):i-1))/g[1]
end
return quotient
end
Whose derivative I have defined as follows:
function ChainRulesCore.rrule(::typeof(factor_divide),poly::Vector{ComplexF64},α::ComplexF64,p::Int64)
### Can this be done with less allocations??
# quotient::Vector{ComplexF64} = factor_divide(poly, α, p)
# function factor_divide_pullback(co_quotient)
quotient_deg::Int64 = size(poly,1) - p - 1
quotient::Vector{ComplexF64} = zeros(ComplexF64, quotient_deg + 1)
g::Vector{ComplexF64} = map(i->(-α)^(p-i+1)*binomial(p,i-1),1:p+1)
g_deriv::Vector{ComplexF64} = map(i->-(p-i+1)*(-α)^(p-i)*binomial(p,i-1),1:p)
co_f = NoTangent()
co_poly::Matrix{ComplexF64} = zeros(ComplexF64,size(poly,1),size(quotient,1)) ## Transpose. Column Major.
co_α::Matrix{ComplexF64} = zeros(ComplexF64,1, size(quotient,1))
co_p = NoTangent()
quotient[1] = poly[1]/g[1]
co_poly[1,1] = 1/g[1]
co_α[1,1] = -quotient[1]*g_deriv[1]/g[1]
@inbounds for i in 2:size(quotient, 1)
quotient[i] = (poly[i]-sum(j->quotient[j]*g[i-j+1],max(1,i-p):i-1))/g[1]
co_α[1,i] = -(quotient[i]*g_deriv[1]+mapreduce(j->co_α[j,1]*g[i-j+1] + quotient[j]*g_deriv[i-j+1],+,max(1,i+1-p):i-1;init=zero(ComplexF64)))/g[1]
@inbounds for j in 1:i
co_poly[j,i] = (I[j,i]-sum(k->co_poly[j,k]*g[i+1-k],max(1,i-p):i-1))/g[1]
end
end
function factor_divide_pullback(co_quotient)
return co_f, co_poly * co_quotient, co_α*co_quotient, co_p
end
return quotient, factor_divide_pullback
end
When I try to take the function’s gradient directly, there seems to be no problem: I use Zygote’s gradient or jacobian.
However, when I use a test case of the following form:
Zygote.gradient((x,y,z)->real(evalpoly(randn(ComplexF64), complex([-1.0,3.0,-1.0,1.0]),complex(1.0),1)
Zygote raises an error:
MethodError: no method matching (::ProjectTo{ComplexF64, NamedTuple{(), Tuple{}}})(::Vector{ComplexF64})
Closest candidates are:
(::ProjectTo{T})(::Integer) where T<:(Complex{<:AbstractFloat})
@ ChainRulesCore ~/.julia/packages/ChainRulesCore/C73ay/src/projection.jl:181
(::ProjectTo{T})(::AbstractZero) where T
@ ChainRulesCore ~/.julia/packages/ChainRulesCore/C73ay/src/projection.jl:120
(::ProjectTo{<:Number})(::Tangent{<:Complex})
@ ChainRulesCore ~/.julia/packages/ChainRulesCore/C73ay/src/projection.jl:192
...
Stacktrace:
[1] _project
@ ~/.julia/packages/Zygote/AS0Go/src/compiler/chainrules.jl:184 [inlined]
[2] map
@ ./tuple.jl:299 [inlined]
[3] map(f::typeof(Zygote._project), t::Tuple{Vector{ComplexF64}, ComplexF64, Int64}, s::Tuple{Vector{ComplexF64}, Vector{ComplexF64}, Nothing})
@ Base ./tuple.jl:302
[4] gradient(::Function, ::Vector{ComplexF64}, ::Vararg{Any})
@ Zygote ~/.julia/packages/Zygote/AS0Go/src/compiler/interface.jl:98
[5] top-level scope
@ In[131]:1
I am, however, able to take a derivative using the first argument. The problem seems to be whenever I also try to include the second argument. Any help would be greatly appreciated.