Make mutating function more AD-friendly

General Usage Performance
, ,
5

I cannot blame you, DI is super handy!

Both of your suggestion threw the very same error:

ERROR: Constant memory is stored (or returned) to a differentiable variable.
As a result, Enzyme cannot provably ensure correctness and throws this error.
This might be due to the use of a constant variable as temporary storage for active memory (https://enzyme.mit.edu/julia/stable/faq/#Runtime-Activity).
If Enzyme should be able to prove this use non-differentable, open an issue!
To work around this issue, either:
 a) rewrite this variable to not be conditionally active (fastest, but requires a code change), or
 b) set the Enzyme mode to turn on runtime activity (e.g. autodiff(set_runtime_activity(Reverse), ...) ). This will maintain correctness, but may slightly reduce performance.
 Failure within method: MethodInstance for PIXIE.deposit!(::Vector{…}, ::SimpleResponse, ::ComponentArrays.ComponentVector{…}, ::Vector{…})
Hint: catch this exception as `err` and call `code_typed(err)` to inspect the errornous code.
If you have Cthulu.jl loaded you can also use `code_typed(err; interactive = true)` to interactively introspect the code.
Mismatched activity for:   store {} addrspace(10)* %3, {} addrspace(10)* addrspace(11)* %.fca.2.gep59, align 8, !dbg !242, !noalias !194 const val: {} addrspace(10)* %3
 value=Unknown object of type Vector{Float64}
 llvalue={} addrspace(10)* %3

Stacktrace:
 [1] tile_halves
   @ ~/.julia/packages/Tullio/2zyFP/src/threads.jl:136
 [2] threader
   @ ~/.julia/packages/Tullio/2zyFP/src/threads.jl:65
 [3] macro expansion
   @ ~/.julia/packages/Tullio/2zyFP/src/macro.jl:1004
 [4] deposit!
   @ ~/Documents/PhD/Alignment/PIXIE/src/covariates_response_models.jl:127

Stacktrace:
  [1] tile_halves
    @ ~/.julia/packages/Tullio/2zyFP/src/threads.jl:136 [inlined]
  [2] threader
    @ ~/.julia/packages/Tullio/2zyFP/src/threads.jl:65 [inlined]
  [3] macro expansion
    @ ~/.julia/packages/Tullio/2zyFP/src/macro.jl:1004 [inlined]
  [4] deposit!
    @ ~/Documents/PhD/Alignment/PIXIE/src/covariates_response_models.jl:127
  [5] npll!
    @ ~/Documents/PhD/Alignment/PIXIE/src/pseudologlikelihood.jl:40
  [6] f
    @ ./REPL[12]:1 [inlined]
  [7] f
    @ ./REPL[12]:0 [inlined]
  [8] diffejulia_f_12219_inner_71wrap
    @ ./REPL[12]:0
  [9] macro expansion
    @ ~/.julia/packages/Enzyme/eJcor/src/compiler.jl:5875 [inlined]
 [10] enzyme_call
    @ ~/.julia/packages/Enzyme/eJcor/src/compiler.jl:5409 [inlined]
 [11] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/eJcor/src/compiler.jl:5295 [inlined]
 [12] autodiff
    @ ~/.julia/packages/Enzyme/eJcor/src/Enzyme.jl:521 [inlined]
 [13] autodiff
    @ ~/.julia/packages/Enzyme/eJcor/src/Enzyme.jl:542 [inlined]
 [14] macro expansion
    @ ~/.julia/packages/Enzyme/eJcor/src/sugar.jl:286 [inlined]
 [15] gradient(::ReverseMode{…}, ::typeof(f), ::ComponentArrays.ComponentVector{…}, ::PIXIE.IsingCacheCPU{…}, ::Const{…}, ::Const{…}, ::Const{…})
    @ Enzyme ~/.julia/packages/Enzyme/eJcor/src/sugar.jl:273
 [16] top-level scope
    @ REPL[18]:1
Some type information was truncated. Use `show(err)` to see complete types.

Looking at my code this is where things fail:

function deposit!(h, R::SimpleResponse, p, x)
    if size(h) ≠ (num_nodes(R),)
        throw(DimensionMismatch("h vector size does not match response model"))
    end
    Tv = eltype(h)
    # Energy deposit response
    hb = p.deposit.bias
    hl = p.deposit.linear
    hq = p.deposit.quadratic
    # Bias part
    h .= hb
    # Adding linear part
    mul!(h, hl, x, one(Tv), one(Tv))
    # Adding quadratic part
    # Element wise quadratic form. Einstein notation
    @tullio h[i] += 0.5 * hq[i, j, k] * x[j] * x[k] # <-- THIS IS THE LINE THAT GIVES ERROR!
    return nothing
end

Is it Tullio known for having issues in integrating with Enzyme?

P.S. Sorry for not providing a MWE, I will try later this week to creating it.