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