Enzyme compilation failed when using logpdf

[BenjaminJCox]

Example:

using Enzyme
using Random
using Distributions
using DistributionsAD

function test_logpdf(X)
    _l = 0.0
    z = zero(X[1, :])
    C = collect(1.0 .* I(length(z)))
    for i in 1:size(X, 1)
        _l = _l + logpdf(MvNormal(X[i, :], C), z)
    end
    return _l
end

ntx = 2.0 .* randn(20, 4)

enz_lmf = Enzyme.gradient(Enzyme.Reverse, θ -> test_logpdf(θ), ntx)

fails compilation

error is long and is given in error.txt

System info:

Julia Version 1.8.4
Commit 00177ebc4f (2022-12-23 21:32 UTC)
Platform Info:
  OS: Windows (x86_64-w64-mingw32)
  CPU: 8 × Intel(R) Core(TM) i7-9700K CPU @ 3.60GHz
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-13.0.1 (ORCJIT, skylake)
  Threads: 7 on 8 virtual cores

apologies if this is not enough information, please let me know if anything more is needed.

[wsmoses]

This looks to be lapack in particular,

declare void @dpotrf_64_(i8*, i8*, i64, i8*, i64, i32) local_unnamed_addr #55

[BenjaminJCox]

This looks to be lapack in particular,

declare void @dpotrf_64_(i8*, i8*, i64, i8*, i64, i32) local_unnamed_addr #55

Is this an issue on my end or in general? Is there anything I could run to better assist in this issue? Thanks!

[wsmoses]

This is an issue in general that we don't currently handle lapack functions.

We (mostly @ZuseZ4 ) actually are presently working on adding BLAS & LaPack support at the moment.

Alternatively, you could write a custom rule for whatever function calls the lapack function. @sethaxen may be able to help with this (see https://enzyme.mit.edu/julia/stable/generated/custom_rule/ for docs on custom rules).

Finally, it may be useful to get a partial backtrace from the location that calls the lapack function, rather than just error if its called anywhere. To do so, you can add the function name (in this case dpotrf_64_ ) to this tuple:

Enzyme.jl/src/compiler.jl

Line 6941 in c807462

(),

[sethaxen]

Alternatively, you could write a custom rule for whatever function calls the lapack function. @sethaxen may be able to help with this (see https://enzyme.mit.edu/julia/stable/generated/custom_rule/ for docs on custom rules).

This is the rule to port to Enzyme:
https://github.com/JuliaDiff/ChainRules.jl/blob/884981745b141f9748aaa5a50cb1ba1a9618c498/src/rulesets/LinearAlgebra/factorization.jl#L476

There are some example implementations of BLAS rules here: #739

[wsmoses]

In any case since this is lapack support, am closing as a dup of #241