[WIP] Add support for Enzyme #1887
Open
+25
−22
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Pull Request Test Coverage Report for Build 3800188054
💛 - Coveralls |
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #1887 +/- ##
======================================
Coverage 0.00% 0.00%
======================================
Files 21 21
Lines 1367 1437 +70
======================================
- Misses 1367 1437 +70 ☔ View full report in Codecov by Sentry. |
devmotion
commented
Dec 23, 2022
devmotion
commented
Dec 23, 2022
|
Also if you want to disable the warnings you can set it like so (https://github.com/EnzymeAD/Enzyme.jl/blob/c29e6119c7963ddb22f1363726f762455748e193/src/api.jl#L414 |
|
You also may want to set the version to 0.11.2 since your CI currently is running at 0.11.0 ( |
|
@devmotion this PR (EnzymeAD/Enzyme.jl#914) should fix the immediate issues you see on CI if you want to try. |
|
This would not be a meaningful definition of |
Looks like the function LinearAlgebra.LAPACK.potrf! is the source of the issue. Is there a description for how to add rules? |
|
Yes we have a section in the docs about Julia level rules and it also links to several examples: https://enzyme.mit.edu/julia/dev/ Alternatively you may consider adding the potrf rule to our internal BLAS handler (which wraps it up for many abis like cblas fortran blas cuBLAS etc): https://github.com/EnzymeAD/Enzyme/blob/2d0b3bbba17abe698a0f15321cb254a2aa189c17/enzyme/Enzyme/BlasDerivatives.td#L151 |
@devmotion @wsmoses do we know where this function call originated from? I don't understand why we have such an operation.
It is very good news to see all segfaults gone, but I agree with @devmotion that more tests are needed to be sure. I think TuringLang/DistributionsAD.jl#254 is better than this PR since Distributions logdensity diff tests cover a lot more standard library functions. |
|
@sunxd3 @wsmodes What functions are you using on the Cholesky? |
I'm happy to take a stab at this over the next few weeks. I've never worked in TableGen though. Any tips before I get started? |
|
@devmotion the infinite recursion is a known issue that comes from Julia 1.10 see EnzymeAD/Enzyme.jl#1134 which was caused by JuliaSparse/SparseArrays.jl@c402d09. Can that test be disabled for now (and ideally SparseArrays.jl gets fixed) |
|
My plan is to update the PR when the |
|
@devmotion the symmetric PR is merged, as well as also removing of the memmove errors. |
|
@devmotion the only two failing tests now are the following. Since you took a look at cholesky earlier, mind investigating the present mismatch? a cholesky dimension mismatch: and the stack overflow caused by a type unstable hvcat hitting the SparseArrays.jl bug I mentioned |
|
I'm busy with other things this weekend but will have a look next week. I'm a bit surprised about the dimension mismatch, I thought this issue was fixed by a few of the improvements to the |
|
A few breaking examples. Truncated normalSimple model with only a This one ends in a segfault State space model with matricesA state space model where specifically the matrix version breaks but the "non-matrix" version doesn't (vector of vectors). If you run the following code: using TuringEnzymeCon2024
using Enzyme
using Turing, LogDensityProblems, LogDensityProblemsAD, ADTypes
Enzyme.API.runtimeActivity!(true);
Enzyme.API.typeWarning!(false);
example_ok = TuringEnzymeCon2024.SatelliteExample()
model_ok = TuringEnzymeCon2024.make_model(example_ok, 10)
f_ok = ADgradient(AutoEnzyme(), Turing.LogDensityFunction(model_ok))
LogDensityProblems.logdensity_and_gradient(f_ok, Turing.VarInfo(model_ok)[:])
example_fail = TuringEnzymeCon2024.SatelliteMatrixExample()
model_fail = TuringEnzymeCon2024.make_model(example_fail, 10)
f_fail = ADgradient(AutoEnzyme(), Turing.LogDensityFunction(model_fail))
LogDensityProblems.logdensity_and_gradient(f_fail, Turing.VarInfo(model_fail)[:])This one results in the matrix being converted into |
@torfjelde, is there an MWE example that only depends on Distributions.jl? |
|
The Cholesky-related test is still failing. |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Note: This does not work yet
I opened this PR to make it easier to debug (and possibly fix) issues with Enzyme.
Currently, the following example does
notwork (note that the snippet does not require the PR which solely reexportsAutoEnzymeat this point):With Enzyme#main my Julia (1.8.1) segfaults. An incomplete (it filled my whole terminal) output: https://gist.github.com/devmotion/1352197f2354c6fecddd7b778ec4bcf7#file-log-txtThe example works (latest releases of Turing, Enzyme, and ADTypes on Julia 1.10.0) but the following warnings show up: