User defined model does not work

[rs1909]

I have defined a model that is a multivariate polynomial. For some reason Flux.jacobian returns zero no matter what. ForwardDiff does not work either.

ForwardDiff gives:
ERROR: LoadError: TypeError: in typeassert, expected Float64, got ForwardDiff.Dual{Nothing,Float64,4}
I believe something is not being tracked, but I have no clue what.

I have posted the issue with a slightly different code here: https://discourse.julialang.org/t/flux-user-defined-layer-problem/20455?u=fastwave

The code is

using Flux
using LinearAlgebra
using DynamicPolynomials
using MultivariatePolynomials

struct PolyModel{S,T}
    mexp::S
    W::T
end

function 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))))
end

function (a::PolyModel)(z::AbstractArray)
    mon = ones(length(a.mexp))
    for j=1:length(a.mexp)
        for k=1:length(z)
            for p=1:a.mexp[j][k]
                mon[j] *= z[k]
            end
        end
    end
    return a.W*mon
end

Flux.@treelike PolyModel

a = PolyModel(2,2,2)
j1 = Flux.jacobian(a,zeros(2))
j2 = ForwardDiff.jacobian(a,zeros(2))

[KristofferC]

See bullet three at http://www.juliadiff.org/ForwardDiff.jl/latest/user/limitations.html.

[rs1909]

Thanks for the suggestion. I could make ForwardDiff work. But Flux.jacobian still does not work

The changes are:

...
function (a::PolyModel)(z::AbstractArray{T}) where T
    mon = ones(T, length(a.mexp))
...

and added
Base.convert(::Type{T}, x::Flux.Tracker.TrackedReal{T}) where T<:Real = x.data

The error now becomes:

ERROR: LoadError: MethodError: no method matching back!(::Float64; once=false)
Closest candidates are:
  back!(::Any, ::Any; once) at /Users/rs1909/.julia/packages/Flux/Dkez4/src/tracker/back.jl:63
  back!(::Flux.Tracker.TrackedReal; once) at /Users/rs1909/.julia/packages/Flux/Dkez4/src/tracker/lib/real.jl:14
  back!(::Flux.Tracker.TrackedArray) at /Users/rs1909/.julia/packages/Flux/Dkez4/src/tracker/lib/array.jl:66 got unsupported keyword argument "once"

[MikeInnes]

The reason that conversion method isn't defined is that it drops gradient information (otherwise we'd just use Float64 everywhere). If you want this to work you'll have to make the ones container Array{Any}.

[dawbarton]

To construct the polynomials in such a way that Flux can deal with them, try

function PolyModel(in::Integer, out::Integer, order::Integer)
    @polyvar x[1:in]
    mx0 = monomials(x, 0:order)
    mx1 = hcat([exponents(m) for m in mx0]...)
    return PolyModel(mx1, param(randn(out, length(mx0))))
end

function (a::PolyModel)(z)
    zp = z.^a.mexp
    zpprod = zp[1, :] # doing prod(zp, dims=1) would be nicer but doesn't work with Flux
    for i = 2:size(zp, 1)
        zpprod = zpprod .* zp[i, :]
    end
    return a.W*zpprod
end

This works with both Flux and ForwardDiff.

[mcabbott]

Glad to see someone else has a use for prod(A, dims=1), you can try with #524 or equivalently this rough package for now.