Thanks, @mcabbott . I must say that I do not really understand the Broadcast function, and I do not (so far) consider it very flexible, probably due to a lack of understanding. I created a MWE (before acting on your last message) of the problem I am now confronted with. I consider a 1D polynomial in x
with a single term: a constant. I still get the same error:
using Zygote
using Tullio
function generate_polynomial(N, degree, x, y)
println("Enter Polynomial::generate_polynomial, length x, y: ", size(x), size(y)) # x: (128, 2), y: (2, 128)
Ix = 0:degree
Jx = 0:degree
# for 2D polynomials
I = [(i,j) for i in Ix for j in Jx if i+j <= degree] # Equivalent to the above loop
# for 1D polynomials
I = [i for i in Ix]
# Construct a polynomial in x and y, according to the indices listed in I
#Poly(coef) = @tullio poly[i] := coef[j] * x[i]^I[j][1] * y[i]^I[j][2] grad=Dual nograd=x nograd=y nograd=I verbose=2
# 1-D polynomial
Poly(coef) = @tullio poly[i] := coef[j] * x[i]^I[j][1] grad=Dual nograd=x nograd=I verbose=2
return Poly #, coef
end
N = 1
degree = 0 # single term
x = rand(N)
y = rand(N)
#poly, coef = generate_polynomial(N, degree, x, y);
poly = generate_polynomial(N, degree, x, y);
nb_terms = Int((degree+1)*(degree+2) / 2)
coef = rand(nb_terms)
loss = x -> sum(poly(x).^2)
Zygote.gradient(loss, coef) # ======= ERROR
Here is the output:
Info: >>>>> Preliminary expressions
β verbosetidy(ex_pre) =
β quote
β (all)((πΆπ->begin
β (axes)(πΆπ) == (axes)((first)(I))
β end), I) || (throw)("elements I[j] must be of uniform size")
β end
β Info: ===== Base actor
β verbosetidy(ex_act) =
β quote
β local @inline(function ππΈπ!(::Type, β::AbstractArray{π―}, coef, x, I, πΆπi, πΆπj, β»οΈ = nothing, π = true) where π―
β @inbounds @fastmath(begin
β for i = πΆπi
β ππΈπΈ = if β»οΈ === nothing
β zero(π―)
β else
β β[i]
β end
β for j = πΆπj
β ππΈπΈ = ππΈπΈ + coef[j] * x[i] ^ (I[j])[1]
β end
β β[i] = ππΈπΈ
β end
β end)
β end)
β end
β Info: ===== Base actor (gradient using ForwardDiff)
β verbosetidy(ex_act) =
β quote
β local @inline(function βππΈπ!(::Type, π₯coef, π₯x, π₯I, π₯β::AbstractArray{π―}, β, coef, x, I, πΆπi, πΆπj, β»οΈ = nothing, π = true) where π―
β @inbounds @fastmath(begin
β begin
β πcoef = ForwardDiff.Dual((zero)(π―), ((one)(π―),))
β end
β begin
β for j = πΆπj
β for i = πΆπi
β β = (coef[j] + πcoef) * x[i] ^ π³
β π₯coef[j] = π₯coef[j] + ForwardDiff.partials(β, 1) * π₯β[i]
β end
β end
β end
β end)
β end)
β end
β Warning: can't parallelise this gradient, no shared indices (gradient using ForwardDiff)
β @ Tullio ~/.julia/packages/Tullio/NGyNM/src/macro.jl:1061
β Info: ===== Base actor (gradient method for FillArrays)
β verbosetidy(ex_act) =
β quote
β local @inline(function βππΈπ!(::Type, π₯coef, π₯x, π₯I, π₯β::Zygote.Fill{π―}, β, coef, x, I, πΆπi, πΆπj, β»οΈ = nothing, π = true) where π―
β @inbounds @fastmath(begin
β begin
β πcoef = ForwardDiff.Dual((zero)(π―), ((one)(π―),))
β π₯β_value = π₯β.value
β end
β begin
β for j = πΆπj
β for i = πΆπi
β β = (coef[j] + πcoef) * x[i] ^ π³
β π₯coef[j] = π₯coef[j] + ForwardDiff.partials(β, 1) * π₯β_value
β end
β end
β end
β end)
β end)
β end
β Warning: can't parallelise this gradient, no shared indices (gradient method for FillArrays)
β @ Tullio ~/.julia/packages/Tullio/NGyNM/src/macro.jl:1061
[ Info: using ForwardDiff gradient
β Info: <<<<< Gradient maker function
β verbosetidy(ex_make) =
β quote
β local function ββ³πΆπβ―(π₯β::AbstractArray{π―}, β, coef, x, I) where π―
β local π₯coef = fill!(similar(coef, Base.promote_type(eltype(coef), π―)), 0)
β local π₯x = nothing
β local π₯I = nothing
β local πΆπj = (Tullio.linearindex)(coef)
β (Tullio.linearindex)(I) == (Tullio.linearindex)(coef) || (throw)("range of index j must agree")
β local πΆπi = (Tullio.linearindex)(x)
β (Tullio.βthreader)(βππΈπ!, (Tullio.storage_type)(π₯coef, π₯x, π₯I, coef, x, I), tuple(π₯coef, π₯x, π₯I, π₯β, β, coef, x, I), tuple(), tuple(πΆπi, πΆπj), 87382)
β return (π₯coef, π₯x, π₯I)
β end
β end
β Info: threading threshold (from cost = 3)
β block = 87382
β Info: >>>>> Maker function
β verbosetidy(ex_make) =
β quote
β local @inline(function β³πΆπβ―(coef, x, I)
β local πΆπj = (Tullio.linearindex)(coef)
β (Tullio.linearindex)(I) == (Tullio.linearindex)(coef) || (throw)("range of index j must agree")
β local πΆπi = (Tullio.linearindex)(x)
β @info "left index ranges" i = (Tullio.linearindex)(x)
β @info "reduction index ranges" j = (Tullio.linearindex)(coef)
β local ππ½π(coef, x, I, j, i) = begin
β identity(coef[j] * x[i] ^ (I[j])[1])
β end
β begin
β local π―1 = Core.Compiler.return_type(ππ½π, (typeof)((coef, x, I, (first)(πΆπj), (first)(πΆπi))))
β local π―2 = if Base.isconcretetype(π―1)
β π―1
β else
β @warn "unable to infer eltype from RHS"
β (typeof)(ππ½π(coef, x, I, (first)(πΆπj), (first)(πΆπi)))
β end
β end
β local π―3 = π―2
β local π― = π―3
β local poly = similar(coef, π―, tuple(πΆπi))
β begin
β (Tullio.threader)(ππΈπ!, (Tullio.storage_type)(poly, coef, x, I), poly, tuple(coef, x, I), tuple(πΆπi), tuple(πΆπj), +, 87382, nothing)
β poly
β end
β end)
β end
β store.
β arrays = [:coef, :x, :I]
β avx = true
β axisdefs = Expr[:(local πΆπj = (Tullio.linearindex)(coef)), :((Tullio.linearindex)(I) == (Tullio.linearindex)(coef) || (throw)("range of index j must agree")), :(local πΆπi = (Tullio.linearindex)(x))]
β constraintpairs = Tuple[]
β constraints = Dict{Symbol, Vector}(:j => Any[:((Tullio.linearindex)(coef)), :((Tullio.linearindex)(I))], :i => Any[:((Tullio.linearindex)(x))])
β cost = 3
β cuda = true
β fastmath = true
β finaliser = :identity
β grad = :Dual
β init = :(zero(π―))
β initkeyword = :π―
β leftarray = :poly
β leftind = [:i]
β leftnames = Symbol[]
β leftraw = Any[:i]
β leftscalar = nothing
β mod = Main
β newarray = true
β nograd = [:x, :I]
β notfree = Symbol[]
β outex = Expr[:(local πΆπj = (Tullio.linearindex)(coef)), :((Tullio.linearindex)(I) == (Tullio.linearindex)(coef) || (throw)("range of index j must agree")), :( ...
β outpre = Expr[:((all)((πΆπ->begin
(axes)(πΆπ) == (axes)((first)(I))
end), I) || (throw)("elements I[j] must be of uniform size")) ...
β padkeyword = :π―
β padmodclamp = false
β plusequals = false
β redfun = :+
β redind = [:j]
β right = :(coef[j] * x[i] ^ (I[j])[1])
β rightind = [:j, :i]
β scalar = false
β scalars = Symbol[]
β sharedind = Symbol[]
β shiftedind = Symbol[]
β threads = true
β unsafeleft = Symbol[]
β unsaferight = Symbol[]
β verbose = 2
β zero = false
β
generate_polynomial
1
0
1-element Vector{Float64}:
0.7905131946859776
1-element Vector{Float64}:
0.9845248220058904
Enter Polynomial::generate_polynomial, length x, y: (1,)(1,)
1
1-element Vector{Float64}:
0.15588748892996085
#190 (generic function with 1 method)
β Info: left index ranges
β i = Base.OneTo(1)
β Info: reduction index ranges
β j = Base.OneTo(1)
ERROR: UndefVarError: π³ not defined
Stacktrace:
[1] βππΈπ!
@ ~/.julia/packages/Tullio/NGyNM/src/macro.jl:1211 [inlined]
[2] tile_halves(fun!::var"#βππΈπ!#186", ::Type{Union{Nothing, Vector{Float64}}}, As::Tuple{Vector{Float64}, Nothing, Nothing, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{Int64}}, Is::Tuple{}, Js::Tuple{UnitRange{Int64}, UnitRange{Int64}}, keep::Nothing, final::Bool)
@ Tullio ~/.julia/packages/Tullio/NGyNM/src/threads.jl:139
[3] tile_halves
@ ~/.julia/packages/Tullio/NGyNM/src/threads.jl:136 [inlined]
[4] βthreader(fun!::var"#βππΈπ!#186", #unused#::Type{Union{Nothing, Vector{Float64}}}, As::Tuple{Vector{Float64}, Nothing, Nothing, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{Float64}, Vector{Int64}}, I0s::Tuple{}, J0s::Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, block::Int64)
@ Tullio ~/.julia/packages/Tullio/NGyNM/src/threads.jl:103
[5] (::var"#942#ββ³πΆπβ―#187"{var"#βππΈπ!#186"})(π₯β::Vector{Float64}, β::Vector{Float64}, coef::Vector{Float64}, x::Vector{Float64}, I::Vector{Int64})
@ Main ~/.julia/packages/Tullio/NGyNM/src/macro.jl:1300
[6] tullio_back
@ ~/.julia/packages/Tullio/NGyNM/src/eval.jl:53 [inlined]
[7] (::Zygote.ZBack{Tullio.var"#tullio_back#156"{Tullio.Eval{var"#β³πΆπβ―#188"{var"#ππΈπ!#185"}, var"#942#ββ³πΆπβ―#187"{var"#βππΈπ!#186"}}, Tuple{Vector{Float64}, Vector{Float64}, Vector{Int64}}, Vector{Float64}}})(dy::Vector{Float64})
@ Zygote ~/.julia/packages/Zygote/SmJK6/src/compiler/chainrules.jl:206
[8] Pullback
@ ~/src/2022/basic_UODE/custom_lux_layer/tullio_zygote_mwe.jl:34 [inlined]
[9] (::typeof(β(Ξ»)))(Ξ::Vector{Float64})
@ Zygote ~/.julia/packages/Zygote/SmJK6/src/compiler/interface2.jl:0
[10] Pullback
@ ~/src/2022/basic_UODE/custom_lux_layer/tullio_zygote_mwe.jl:48 [inlined]
[11] (::typeof(β(#190)))(Ξ::Float64)
@ Zygote ~/.julia/packages/Zygote/SmJK6/src/compiler/interface2.jl:0
[12] (::Zygote.var"#60#61"{typeof(β(#190))})(Ξ::Float64)
@ Zygote ~/.julia/packages/Zygote/SmJK6/src/compiler/interface.jl:45
[13] gradient(f::Function, args::Vector{Float64})
@ Zygote ~/.julia/packages/Zygote/SmJK6/src/compiler/interface.jl:97
[14] top-level scope
@ ~/src/2022/basic_UODE/custom_lux_layer/tullio_zygote_mwe.jl:49
One can see the appearance of the \teapot
symbol during the derivative calculation without it having been defined. Where is comes from I have no idea! (@ChrisRackauckas)