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