Forward and reverse Enzyme tests and rules for linalg#449
Conversation
|
Your PR no longer requires formatting changes. Thank you for your contribution! |
|
The test on 1.12 is passing locally for me! I assume it's getting OOMed or something... |
|
OK, everything looks happy now except the GPU stuff which is unrelated. Are we good to go? |
|
Do we think the test failure is a problem with how LRU interacts with Enzyme? From the stacktrace, I seem to read this as not finding a key, even though being in an if-clause that explicitly checks this: https://github.com/JuliaCollections/LRUCache.jl/blob/1dad9fef75fef51ea1b7e984e5850ad4e374a7e0/src/LRUCache.jl#L172-L175 The really confusing part to me is that it seems to originate from a forward call, which should just be a regular function call, so I'm not sure what is really going on there. I also don't think this can really be a race condition, since 1) I don't think we are multithreading, 2) LRU protects against this? |
|
It also seems to only happen in the CompatCheck tests, not the main ones |
|
Let me just see if bumping the LRUCache compat helps at all... |
|
OK that makes CompatCheck pass and the min test fail. It seems the failures only happen on 1.10 regardless, but they are intermittent. Really annoying. |
|
Also locally I can see it happen in reverse calls so I think it's not to do with fwd mode really |
|
Removing the |
|
OK so after more digging, it looks like the problem here is 1.10 + Enzyme + the |
| ΔAr = NoRData() | ||
|
|
||
| Δαr = isnothing(Ap) ? NoRData() : project_scalar(α, inner(Ap, ΔC)) | ||
| Δαr = isnothing(Ap) ? NoRData() : TO.project_scalar(α, inner(Ap, ΔC)) |
There was a problem hiding this comment.
I lost track of where our new project_ methods live. Is TO a logical place for project_scalar? We probably also have it in VI? Or is TO just importing the VI version? Where is project_mul!, is that in TO? Or TK?
There was a problem hiding this comment.
Lukas suggested pointing to TO. We have duplicated versions of these in both Mooncake and Enzyme extensions over at VI 🫠
There was a problem hiding this comment.
project_mul lives here at TK
There was a problem hiding this comment.
@Jutho did you ever read those books "Where's Waldo?"? We need that for ad helpers
There was a problem hiding this comment.
I think I originally just wrote these as small helper functions, but then it turned out we kind of need them everywhere, and then they just kind of proliferated
There was a problem hiding this comment.
Can you read "Where's Waldo"? Maybe we can just create one super package VIMAKTOTK.jl, and this solves all of our problems. And then we can keep up with the QMC guys in terms of acronyms 😄 .
There was a problem hiding this comment.
All depends on your definition of "reading" 😉
|
Some more strangeness: the missing key error only ever occurs with |
Trying to make these a little more manageable and pick up the fwd rules where possible