I am trying to define a layer that is a polynomial in Flux. The definition seems to work, but any training with it diverges. It is also impossible to calculate the Jacobian of this polynomial function using Flux, which should be trivial. The following code illustrates the problem with the Jacobian. The question is: what am I doing wrong?

using Flux

using DynamicPolynomials

using MultivariatePolynomialsstruct PolyModel

mexp

W

endfunction PolyModel(in::Integer, out::Integer, order::Integer)

@polyvar x[1:in]

mx0 = monomials(x, 0:order)

mx1 = [exponents(mx0[k]) for k=1:length(mx0)]

return PolyModel(mx1, param(randn(out, length(mx0))))

endFlux.@treelike PolyModel

function (a::PolyModel)(z::AbstractArray)

return a.W*[prod([z[k]^(a.mexp[j][k]) for k=1:length(z)]) for j=1:length(a.mexp)]

end

a = PolyModel(4,4,2)

Flux.jacobian(a,zeros(4))

ERROR: MethodError: no method matching Float64(::Flux.Tracker.TrackedReal{Float64})

Closest candidates are:

Float64(::Real, ::RoundingMode) where T<:AbstractFloat at rounding.jl:185

Float64(::T<:Number) where T<:Number at boot.jl:725

Float64(::Int8) at float.jl:60

…