Return type of function is not a valid datatype?

[wsmoses]

trafficstars

using Enzyme, SparseArrays, LinearAlgebra

# setup data
θ = rand(3)
x = rand(11)
ν = rand(10)
λ = rand(24)
S = sprand(Float64, 6, 7, 0.1)
Cp = sprand(Float64, 1, 11, 0.4)
_h = rand(24)
_d = spzeros(10)
_c = spzeros(11)
λ = zeros(24)
_Q = spzeros(11, 11)
_c[7] = -1.0
θ = [30.0, 70.0, 90.0]
x = [1.9873817034700316, 3.9747634069400632, 1.9873817034700316, 1.9873817034700316, 1.9873817034700316, 3.9747634069400632, 1.9873817034700316, 0.028391167192429023, 0.022082018927444796, 0.13249211356466878, 0.13249211356466878]
ν = [ 0.0, 0.0, 0.028391167192429023, 0.24290220820189273,0.7287066246056783, 0.0, 1.9873817034700316, 3.9747634069400632, 5.962145110410095, 7.9495268138801265,]
λ[end] = -1.9873817034700316
S = [
    1.0  0.0  0.0  0.0  0.0  0.0  -1.0
    0.0  1.0  0.0  0.0  0.0  0.0  -2.0
    0.0  0.0  1.0  0.0  0.0  0.0  -1.0
    0.0  0.0  0.0  1.0  0.0  0.0  -1.0
    0.0  0.0  0.0  0.0  1.0  0.0  -1.0
    0.0  0.0  0.0  0.0  0.0  1.0  -2.0
]
se_vars = ((2, -1.0), (3, -1.0),(3, -3.0), (1, -2.0), (1, -1.0),)

# create functions
Se(θ) = sparse(
    [1, 2, 3, 4, 3],
    [3, 4, 4, 5, 5],
    [k / θ[idx] for (idx, k) in se_vars],
    4,
    7,
)

E(θ) = [
    S spzeros(6, 4)
    Se(θ) I(4)
]

d = _ -> _d

c = _ -> _c

M = _ -> [
    -I(11)
    I(11)
    -Cp
    Cp
]

h = _ -> _h

Q = _ -> _Q

xidxs = 1:length(x)
νidxs = last(xidxs) .+ (1:length(ν))
λidxs = last(νidxs) .+ (1:length(λ))
θidxs = last(λidxs) .+ (1:length(θ))

sparse_F(z) = [
    Q(z[θidxs]) * z[xidxs] + c(z[θidxs]) 
    #E(z[θidxs])' * z[νidxs] - M(z[θidxs])' * z[λidxs]
    #E(z[θidxs]) * z[xidxs] - d(z[θidxs])
    #spdiagm(z[λidxs]) * (M(z[θidxs]) * z[xidxs] - h(z[θidxs]))
]

z = [x; ν; λ; θ]
sparse_F(z) # this function should be differentiated wrt z

function f(out, z)
    out .= sparse_F(z)    
    nothing # the docs say this you should do this for array output funcs
end

dz = zero(z)
@show z, dz
@show fwddiff(sparse_F, Duplicated(z, dz))
@show z, fz

@show sparse_F(z)
@show length(z)
@show fwdjacobian(sparse_F, z; Val(1))
@show fwdjacobian(sparse_F, z; Val(length(z)))
using BenchmarkTools

@btime fwdjacobian(sparse_F, z, Val(1))
@btime fwdjacobian(sparse_F, z, Val(length(z)))

# out = spzeros(45)
# d_out = spzeros(45)
# d_z = zeros(length(z))


# autodiff(f, Const, Duplicated(out, d_out), Duplicated(z, d_z))

Can cause:

ERROR: LoadError: AssertionError: rrt isa DataType
Stacktrace:
 [1] #s397#92
   @ /mnt/Data/git/Enzyme.jl/src/compiler.jl:4080 [inlined]
 [2] var"#s397#92"(F::Any, DF::Any, A::Any, TT::Any, Mode::Any, ModifiedBetween::Any, width::Any, specid::Any, ::Any, f::Any, df::Any, #unused#::Type, tt::Any, #unused#::Type, #unused#::Type, #unused#::Type, #unused#::Any)
   @ Enzyme.Compiler ./none:0
 [3] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any})
   @ Core ./boot.jl:580
 [4] thunk
   @ /mnt/Data/git/Enzyme.jl/src/compiler.jl:4117 [inlined]
 [5] thunk
   @ /mnt/Data/git/Enzyme.jl/src/compiler.jl:4110 [inlined]
 [6] fwddiff
   @ /mnt/Data/git/Enzyme.jl/src/Enzyme.jl:356 [inlined]
 [7] fwddiff(f::typeof(sparse_F), args::Duplicated{Vector{Float64}})
   @ Enzyme /mnt/Data/git/Enzyme.jl/src/Enzyme.jl:381
 [8] top-level scope
   @ show.jl:1047
in expression starting at /mnt/Data/git/Enzyme.jl/sparse.jl:81

As

julia> typeof(sparse_F(z))
Vector{SparseVector{Float64, Int64}} (alias for Array{SparseVector{Float64, Int64}, 1})

julia> Core.Compiler.return_type(sparse_F, Tuple{typeof(z)})
Vector (alias for Array{_A, 1} where _A)

[wsmoses]

Current error

┌ Warning: Returned rooting not fully handled, segfault likely
└ @ Enzyme.Compiler ~/.julia/packages/GPUCompiler/N98un/src/utils.jl:35
ERROR: LoadError: UndefRefError: access to undefined reference
Stacktrace:
 [1] getproperty
   @ ./Base.jl:42 [inlined]
 [2] getindex
   @ /mnt/Data/git/Enzyme.jl/src/compiler.jl:4482 [inlined]
 [3] macro expansion
   @ /mnt/Data/git/Enzyme.jl/src/compiler.jl:4759 [inlined]
 [4] enzyme_call
   @ /mnt/Data/git/Enzyme.jl/src/compiler.jl:4532 [inlined]
 [5] ForwardModeThunk
   @ /mnt/Data/git/Enzyme.jl/src/compiler.jl:4517 [inlined]
 [6] autodiff
   @ /mnt/Data/git/Enzyme.jl/src/Enzyme.jl:398 [inlined]
 [7] autodiff
   @ /mnt/Data/git/Enzyme.jl/src/Enzyme.jl:327 [inlined]
 [8] fwddiff(f::typeof(sparse_F), args::Duplicated{Vector{Float64}})
   @ Enzyme ./deprecated.jl:72
 [9] top-level scope
   @ show.jl:1047