DiffCache/GeneralLazyBufferCache for an array with mixed types

[lazarusA]

Is there a way to make this work? The important part is the preallocation of the mixed tuples in the array.

• DiffCache doesn't work, and
• GeneralLazyBufferCache [don't know if it is the right approach instead]

using ForwardDiff
using PreallocationTools
using StaticArrays

ar_alloc = [(; a=1.0, b=2.0, c=SVector(0.0)), (; a=1.0, b=2.0, c=SVector(0.0))]

function fxy(x, ar_alloc, y)
    for i in 1:2
        a = x + 1
        b = x * x + 2
        c = SVector(x * x * x + 3 + a)
        ar_alloc[i] = (; a, b, c)
    end
    ŷ = [ar_alloc[1].a + ar_alloc[1].c[1], ar_alloc[2].a + +ar_alloc[2].c[1]]
    return sum(abs.(ŷ .- y))
end

y = [2.3, 5.0]
loss(x) = fxy(x, ar_alloc, y)
loss(0.5)

ForwardDiff.gradient(loss, [0.2, 0.5, 0.6])

[lazarusA]

ok. Fair. No stack trace. This is the one for the intended pre-allocated array.

ar_alloc = [(; a=1.0, b=2.0, c=SVector(0.0)), (; a=1.0, b=2.0, c=SVector(0.0))]
ar_alloc = DiffCache(ar_alloc)

ERROR: MethodError: no method matching zero(::Type{NamedTuple{(:a, :b, :c), Tuple{Float64, Float64, SVector{1, Float64}}}})

Closest candidates are:
  zero(::Union{Type{P}, P}) where P<:Dates.Period
   @ Dates ~/.julia/juliaup/julia-1.9.1+0.aarch64.apple.darwin14/share/julia/stdlib/v1.9/Dates/src/periods.jl:51
  zero(::Union{String, Type{String}})
   @ Zarr ~/.julia/packages/Zarr/jgFcc/src/metadata.jl:48
  zero(::Union{Type{<:Zarr.DateTime64}, Zarr.DateTime64})
   @ Zarr ~/.julia/packages/Zarr/jgFcc/src/metadata.jl:47
  ...

Stacktrace:
 [1] zeros(#unused#::Type{NamedTuple{(:a, :b, :c), Tuple{Float64, Float64, SVector{1, Float64}}}}, dims::Tuple{Int64})
   @ Base ./array.jl:585
 [2] zeros(#unused#::Type{NamedTuple{(:a, :b, :c), Tuple{Float64, Float64, SVector{1, Float64}}}}, dims::Int64)
   @ Base ./array.jl:580
 [3] DiffCache(u::Vector{NamedTuple{(:a, :b, :c), Tuple{Float64, Float64, SVector{1, Float64}}}}, siz::Tuple{Int64}, chunk_sizes::Vector{Int64})
   @ PreallocationTools ~/.julia/packages/PreallocationTools/nhCNl/src/PreallocationTools.jl:79
 [4] DiffCache(u::Vector{NamedTuple{(:a, :b, :c), Tuple{Float64, Float64, SVector{1, Float64}}}}, N::Int64; levels::Int64)
   @ PreallocationTools ~/.julia/packages/PreallocationTools/nhCNl/src/PreallocationTools.jl:97
 [5] DiffCache(u::Vector{NamedTuple{(:a, :b, :c), Tuple{Float64, Float64, SVector{1, Float64}}}}, N::Int64)
   @ PreallocationTools ~/.julia/packages/PreallocationTools/nhCNl/src/PreallocationTools.jl:95
 [6] DiffCache(u::Vector{NamedTuple{(:a, :b, :c), Tuple{Float64, Float64, SVector{1, Float64}}}})
   @ PreallocationTools ~/.julia/packages/PreallocationTools/nhCNl/src/PreallocationTools.jl:95
 [7] top-level scope

and for GLBC

ar_allocG = GeneralLazyBufferCache(function (p)
    ar_alloc
end
)
# no error, but then how is this suppose to work in the next steps?
loss(x) = fxy(x, ar_allocG, y)
ForwardDiff.gradient(loss, [0.2, 0.5, 0.6])

ERROR: MethodError: no method matching +(::Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(loss), Float64}, Float64, 3}}, ::Int64)
For element-wise addition, use broadcasting with dot syntax: array .+ scalar

Closest candidates are:
  +(::Any, ::Any, !Matched::Any, !Matched::Any...)
   @ Base operators.jl:578
  +(!Matched::T, ::T) where T<:Union{Int128, Int16, Int32, Int64, Int8, UInt128, UInt16, UInt32, UInt64, UInt8}
   @ Base int.jl:87
  +(!Matched::Union{MathOptInterface.ScalarAffineFunction{T}, MathOptInterface.ScalarQuadraticFunction{T}}, ::T) where T
   @ MathOptInterface ~/.julia/packages/MathOptInterface/BlCD1/src/Utilities/functions.jl:1908
  ...

[ChrisRackauckas]

GLBC the code is just incorrect. You'd have to do something like:

function fxy(x, _ar_alloc, y)
   ar_alloc = _ar_alloc[x]
    for i in 1:2
        a = x + 1
        b = x * x + 2
        c = SVector(x * x * x + 3 + a)
        ar_alloc[i] = (; a, b, c)
    end
    ŷ = [ar_alloc[1].a + ar_alloc[1].c[1], ar_alloc[2].a + +ar_alloc[2].c[1]]
    return sum(abs.(ŷ .- y))
end

and then you need an appropriate allocation:

ar_allocG = GeneralLazyBufferCache(function (p)
    [(; a=p, b=p, c=SVector(p)), (; a=p, b=p, c=SVector(p))]
end
)

and you should be good.

[lazarusA]

Thanks, unfortunately, it seems like is not that easy. After following some error suggestions from your code sample, this is the updated example:

# after some type-size-fixes and using `dot` everywhere[for some reason] (suggested by the errors),
# this is a more representative example
using ForwardDiff
using PreallocationTools: DiffCache, GeneralLazyBufferCache
using StaticArrays

function fxy(x, _ar_alloc, y)
    ar_alloc = _ar_alloc[x]
    for i in 1:2
        a = x .+ 1 # things will start to fail from here without the `dot`, why?, there is no need.
        b = x .* x .+ 2
        c = SVector{2}(x .* x .* x .+ 3 .+ a, x .* x)
        ar_alloc[i] = (; a, b, c)
    end
    ŷ = [ar_alloc[1].a .+ ar_alloc[1].c[1], ar_alloc[2].a .+ ar_alloc[2].c[2]]
    return sum(abs2.(ŷ .- y)) # ForwardDiff.gradient will fail here.
end

ar_allocG = GeneralLazyBufferCache(function (p)
    [(; a=p, b=p, c=SVector{2}(p, p)), (; a=p, b=p, c=SVector{2}(p, p))]
end
)

# do we have a number?
y = [2.3, 5.0]
loss(x) = fxy(x, ar_allocG, y)
loss(0.5) # yes, it outputs a number

25.193125000000002

# the actual test
# but here, this one still fails.
ForwardDiff.gradient(loss, [0.2, 0.5, 0.6])

ERROR: MethodError: no method matching -(::Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(loss), Float64}, Float64, 3}}, ::Float64)
For element-wise subtraction, use broadcasting with dot syntax: array .- scalar

Closest candidates are:
  -(::T, ::T) where T<:Union{Float16, Float32, Float64}
   @ Base float.jl:409
  -(::ForwardDiff.Dual{Tx}, ::AbstractFloat) where Tx
   @ ForwardDiff ~/.julia/packages/ForwardDiff/vXysl/src/dual.jl:144
  -(::ForwardDiff.Dual{Tx}, ::Real) where Tx
   @ ForwardDiff ~/.julia/packages/ForwardDiff/vXysl/src/dual.jl:144
  ...

Stacktrace:
  [1] _broadcast_getindex_evalf
    @ ./broadcast.jl:683 [inlined]
  [2] _broadcast_getindex
    @ ./broadcast.jl:656 [inlined]
  [3] _getindex
    @ ./broadcast.jl:680 [inlined]
  [4] _broadcast_getindex
    @ ./broadcast.jl:655 [inlined]
  [5] getindex
    @ ./broadcast.jl:610 [inlined]
  [6] copy(bc::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Tuple{Base.OneTo{Int64}}, typeof(abs2), Tuple{Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{1}, Nothing, typeof(-), Tuple{Vector{Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(loss), Float64}, Float64, 3}}}, Vector{Float64}}}}})
    @ Base.Broadcast ./broadcast.jl:912
  [7] materialize
    @ ./broadcast.jl:873 [inlined]
  [8] fxy(x::Vector{ForwardDiff.Dual{ForwardDiff.Tag{typeof(loss), Float64}, Float64, 3}}, _ar_alloc::GeneralLazyBufferCache{var"#9#10"}, y::Vector{Float64})

[thomvet]

Not sure what you tried doing here, but since loss([0.2, 0.5, 0.6]) gives a ERROR: MethodError: no method matching -(::Vector{Float64}, ::Float64), I'd say that your function is already part of the problem here, because it is not doing what you think it is doing.

A few comments:

1. You tested whether the function gives a number by feeding it a single number...it does, but when you feed it a n x 1 vector, your ŷ becomes a vector of two n x 1 vectors and then you are trying to substract a n x 1 vector from that... unclear what you are trying to do there.

2. SVectors are immutable, so caching them like this (which implies that you want to mutate them later on) makes little sense.

3. The fact that you had to use the dots makes sense, because you are doing the equivalent of [1, 2, 3] .+ 1, i.e., adding a number to a vector in your a = x .+ 1 # things will start to fail from here without the dot, why?, there is no need. line...

I don't think this is an actual bug in the package.