Segfault or silently wrong result on differentiation of GP loss function

[simsurace]

This came up while investigating this topic on Discourse. The following script iterates over a few backends and prints gradients. Depending on the size of N (and luck), this will either give different results for the Enzyme backends vs. FiniteDiff or Zygote or even segfault. When it segfaulted, I got the following error:

[74750] signal (11.2): Segmentation fault: 11
in expression starting at REPL[19]:1
.Ldgemm_tcopy_L8_M8_80 at ~/.julia/juliaup/julia-1.10.9+0.aarch64.apple.darwin14/lib/julia/libopenblas64_.dylib (unknown line)
Allocations: 87827445 (Pool: 87719349; Big: 108096); GC: 101
[1]    74750 segmentation fault  julia +lts --project=. -t4

using AbstractGPs, DifferentiationInterface, Optim, StatsFuns, Random

const DI = DifferentiationInterface

import Enzyme
import FiniteDiff
import Zygote

rng = MersenneTwister(1)

N = 1000
x = randn(rng, N)
y = randn(rng, N)

function loss_function(x, y)
    function negativelogmarginallikelihood(params)
        kernel =
            softplus(params[1]) * (Matern52Kernel() ∘ ScaleTransform(softplus(params[2])))
        noise_var = 0.1
        f = GP(kernel)
        fx = f(x, noise_var)
        return -logpdf(fx, y)
    end
    return negativelogmarginallikelihood
end

θ0 = randn(rng, 2)

losses = [
    loss_function(x[1:4], y[1:4]),
    loss_function(x, y)
]

backends = [
    AutoFiniteDiff(),
    AutoZygote(),
    AutoEnzyme(
        mode=Enzyme.Reverse,
        function_annotation=Enzyme.Duplicated,
    ),
    AutoEnzyme(
        mode=Enzyme.set_runtime_activity(Enzyme.Reverse),
        function_annotation=Enzyme.Const,
    ),
]

for loss in losses
    for backend in backends
        grad = DI.gradient(loss, backend, θ0)
        @info "Gradient" backend grad
    end
end

julia> versioninfo()
Julia Version 1.10.9
Commit 5595d20a287 (2025-03-10 12:51 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin24.0.0)
  CPU: 12 × Apple M3 Pro
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, apple-m1)
Threads: 4 default, 0 interactive, 2 GC (on 6 virtual cores)
Environment:
  JULIA_EDITOR = code

[wsmoses]

Thanks for the report!

This input code is rather complex, can you try to simplify it (ideally being as close to pure julia as possible, using as few as possible external packages as possible), and directly calling Enzyme.autodiff?

This will allow us to find and fix the error fastest

[simsurace]

This is essentially a Gaussian likelihood with a parameterized covariance matrix. There's a cholesky factorization and some linalg. I'll try to minimize it.

[simsurace]

This is a MRE for the wrong gradient. No segfaults with this one. The real code above makes a call to a function from Distances.jl with @inbounds, so that could be the problem.

using LinearAlgebra
import Enzyme
import FiniteDiff

function loss_function(x)
    function testfunction(params)
        xs = params[1] * x
        C_mat = zeros(length(xs), length(xs))
        mul!(C_mat, xs, xs')
        return sum(C_mat)
    end
    return testfunction
end

θ0 = randn(1)
x = randn(10_000)
loss = loss_function(x)

enzyme_grad = let dθ = Enzyme.make_zero(θ0)
    Enzyme.autodiff(
        Enzyme.Reverse,
        Enzyme.Duplicated(loss, Enzyme.make_zero(loss)),
        Enzyme.Duplicated(θ0, dθ),
    )
    dθ
end

finitediff_grad = FiniteDiff.finite_difference_gradient(loss, θ0)

enzyme_grad ≈ finitediff_grad # false

EDIT: removed dependency on DifferentiationInterface.jl

[wsmoses]

@simsurace just to confirm does this still err (a lot of relevant fixes have gone in since)

[simsurace]

I just checked the last example even still produces very different results compared to FiniteDiff:

julia> enzyme_grad = let dθ = Enzyme.make_zero(θ0)
           Enzyme.autodiff(
               Enzyme.Reverse,
               Enzyme.Duplicated(loss, Enzyme.make_zero(loss)),
               Enzyme.Duplicated(θ0, dθ),
           )
           dθ
       end
 ** On entry to DSYMM  parameter number  9 had an illegal value
1-element Vector{Float64}:
 -15147.222392413192

julia> finitediff_grad = FiniteDiff.finite_difference_gradient(loss, θ0)
1-element Vector{Float64}:
 -2455.863678399372

julia> enzyme_grad ≈ finitediff_grad # false
false

(jl_swi6WU) pkg> st
Status `/private/var/folders/31/r_tx2_5d4hn30gf8qbb0hjjc0000gp/T/jl_swi6WU/Project.toml`
  [7da242da] Enzyme v0.13.76
  [6a86dc24] FiniteDiff v2.28.1
  [37e2e46d] LinearAlgebra v1.11.0

julia> versioninfo()
Julia Version 1.11.6
Commit 9615af0f269 (2025-07-09 12:58 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin24.0.0)
  CPU: 12 × Apple M3 Pro
  WORD_SIZE: 64
  LLVM: libLLVM-16.0.6 (ORCJIT, apple-m2)
Threads: 1 default, 0 interactive, 1 GC (on 6 virtual cores)

[simsurace]

Wouldn't this have to be covered by unit tests? It's just a mul! and a sum.