Enzyme crashes when using Duplicated(f) for function that modifies argument.

[ajinkya-k]

I am trying to use reverse mode autodiff to "refit" a linear mixed effect model using MixedModels.jl. To that end I wrote a function called adfit! (see code chunk below). Because the objective function takes in a model object and mutates in the course of computing the objective I used Duplicated(fn) based on a suggestion by @wsmoses in slack.

However, that is causing Julia to crash with the following message:

[19562] signal 6: Abort trap: 6
in expression starting at none:0
__pthread_kill at /usr/lib/system/libsystem_kernel.dylib (unknown line)
Allocations: 139784462 (Pool: 139781010; Big: 3452); GC: 52
zsh: abort      julia --proj

Here is the code I am using:

using Revise
using MixedModels
using Enzyme


function adfit!(m::LinearMixedModel)
    f(x) = objective!(m, x)
    x = similar(m.optsum.initial)
    copyto!(x, m.optsum.initial)
    dx = zeros(eltype(x), length(x))
    return autodiff(Reverse, Duplicated(f, make_zero(f)), Active, Duplicated(x, dx))
end

using DataFrames

df = DataFrame(
        grp1 = ["a", "a", "a", "a", "b", "b", "b", "c", "c"],
        grp2 = ["A", "B", "A", "B", "A", "B", "A", "B", "A"],
        x = collect(1.:9.),
        y = randn(9)
)

lmm = fit(LinearMixedModel, @formula(y ~ x + (1 | grp1)), df);

lmm2 = adfit!(lmm)

Here is my julia session info:

julia> VERSION
v"1.11.3"

(testngdir) pkg> st
Status `~/projects/julia_repos/testngdir/Project.toml`
  [13e28ba4] AppleAccelerate v0.4.0
  [0ca39b1e] Chairmarks v1.3.1
  [a93c6f00] DataFrames v1.7.0
  [31a5f54b] Debugger v0.7.10
  [b4f34e82] Distances v0.10.12
  [7da242da] Enzyme v0.13.28
  [28b8d3ca] GR v0.73.12
  [7073ff75] IJulia v1.26.0
  [5903a43b] Infiltrator v1.8.3
  [033835bb] JLD2 v0.5.11
  [98e50ef6] JuliaFormatter v1.0.62
  [2b0e0bc5] LanguageServer v4.5.1
  [b4fcebef] Lasso v0.7.2
  [898213cb] LowRankApprox v0.5.5 `../LowRankApprox.jl`
  [ee78f7c6] Makie v0.22.1
  [ff71e718] MixedModels v4.30.0 `../MixedModels.jl`
  [5fb14364] OhMyREPL v0.5.28
  [91a5bcdd] Plots v1.40.9
  [aafa11c5] PrettyChairmarks v1.0.0 `../PrettyChairmarks.jl#ahk/histcommonaxis`
  [92933f4c] ProgressMeter v1.10.2
  [4c0109c6] QuartoNotebookRunner v0.12.0 `../QuartoNotebookRunner.jl`
⌃ [295af30f] Revise v3.7.1
  [00fa7acb] Statistics v1.11.2 `../Statistics.jl`
  [b8865327] UnicodePlots v3.7.2
  [51474c39] LAPACK_jll v3.12.0+3
  [37e2e46d] LinearAlgebra v1.11.0
  [4536629a] OpenBLAS_jll v0.3.27+1
Info Packages marked with ⌃ have new versions available and may be upgradable.

[ajinkya-k]

Tagging @dmbates here as this might be of interest and he might have some suggestions and pointers