Hi there,
I am struggling to get Enzyme working on a simplified version of my simulation. Here I give the code I’m running along with the error.
My data is stored in an array of Student structs:
using Enzyme
using SpecialFunctions
cdf_eval(x, mean, stddev) = 0.5 * (1 + erf((x - mean) / (stddev * sqrt(2.0))))
struct Student{T}
θ::T
σ::T
u1::T
u2::T
end
The objective function uses the following intermediate function that “simulates” the choice of a student given a cutoff vector (I will be vague on the explanation since this does not matter).
function student_choice!(cutoffs, student, prob)
"""
simple function to simulate the school choice a student given a cutoff vector
and compute its expected value. function modified prob in-place for efficiency motives
"""
# unpack student
(; θ, σ, u1, u2) = student
# we constraint c2 > c1
if u2 ≥ u1
prob[2] = 1.0 - cdf_eval(cutoffs[2], θ, σ)
prob[1] = prob[2] - cdf_eval(cutoffs[1], θ, σ)
expected_u = prob[2] * u2 + prob[1] * u1
else
prob[2] = 0.0
prob[1] = 1.0 - cdf_eval(cutoffs[1], θ, σ)
expected_u = prob[1] * u1
end
return expected_u
end
Then my objective function, maps a cutoff vector to the sum of expected utilities over students in my data. I pass the data inside a parameter vector p since ultimately I want this to be in the form required by Optimization.jl
function f(cutoffs, p)
array_students = p[1]::Vector{Student{Float64}}
probs = zeros(2)
sum = zero(eltype(cutoffs))
@inbounds for student ∈ array_students
probs .= zero(eltype(probs))
sum += student_choice!(cutoffs, student, probs)
end
return sum
end
Note: having this thing with the probs vector might seem weird for now but there’s a reason, not apparent in this simplified MWE.
Then I just simulate data and see whether I can differentiate f
function main()
## SIMULATE DATA ##
number_students = 1000
array_students = Array{Student{Float64}}(undef, number_students)
for student_id ∈ 1:1000
θ = rand()
σ = 0.2
u1 = rand() + 1.0
u2 = rand() + 1.0
array_students[student_id] = Student(θ, σ, u1, u2)
end
## PERFORM OPTIMIZATION ##
# first trivial case
x0 = rand(2)
dcutoffs = zeros(2)
p = [array_students]
autodiff(Reverse, f, Duplicated(x0, dcutoffs), Const(p))
end
main()
But it gives me the following error:
julia> main()
ERROR:
No augmented forward pass found for jl_alloc_genericmemory
at context: %52 = call "enzyme_type"="{[-1]:Pointer}" {} addrspace(10)* @jl_alloc_genericmemory({} addrspace(10)* noundef addrspacecast ({}* inttoptr (i64 4696619088 to {}*) to {} addrspace(10)*), i64 noundef 2) #28, !dbg !126
Stacktrace:
[1] GenericMemory
@ ./boot.jl:516
[2] Array
@ ./boot.jl:578
[3] Array
@ ./boot.jl:591
[4] zeros
@ ./array.jl:589
[5] zeros
@ ./array.jl:586
[6] zeros
@ ./array.jl:584
[7] f
@ ./REPL[21]:3
Stacktrace:
[1] GenericMemory
@ ./boot.jl:516 [inlined]
[2] Array
@ ./boot.jl:578 [inlined]
[3] Array
@ ./boot.jl:591 [inlined]
[4] zeros
@ ./array.jl:589 [inlined]
[5] zeros
@ ./array.jl:586 [inlined]
[6] zeros
@ ./array.jl:584 [inlined]
[7] f
@ ./REPL[21]:3 [inlined]
[8] diffejulia_f_36206wrap
@ ./REPL[21]:0
[9] macro expansion
@ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:7187 [inlined]
[10] enzyme_call
@ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6794 [inlined]
[11] CombinedAdjointThunk
@ ~/.julia/packages/Enzyme/TiboG/src/compiler.jl:6671 [inlined]
[12] autodiff
@ ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:320 [inlined]
[13] autodiff
@ ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:348 [inlined]
[14] autodiff
@ ~/.julia/packages/Enzyme/TiboG/src/Enzyme.jl:329 [inlined]
[15] main()
@ Main ./REPL[45]:18
[16] top-level scope
@ REPL[46]:1
I have made Enzyme work before in similar but even more complicated functions. Can someone help out with pointing out whether I’m missing something here?
Thanks in advance!