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? |
| !isa(A, Const) && TK.project_mul!(C.dval, A.dval, B.val, α.val, One()) | ||
| !isa(B, Const) && TK.project_mul!(C.dval, A.val, B.dval, α.val, One()) | ||
| end | ||
| mul!(C.val, A.val, B.val, α.val, β.val) |
There was a problem hiding this comment.
I guess another case could be made here for computing AB = A.val * B.val separately, in the case where !isa(α, Const). Also, I don't think we need project_mul! here, since C should be complex as soon as any of A, B, α or β is complex.
Maybe something like this:
| mul!(C.val, A.val, B.val, α.val, β.val) | |
| if !isa(C, Const) | |
| scale!(C.dval, β.val) | |
| !isa(β, Const) && add!(C.dval, C.val, β.dval) | |
| !isa(A, Const) && mul!(C.dval, A.dval, B.val, α.val, One()) | |
| !isa(B, Const) && mul!(C.dval, A.val, B.dval, α.val, One()) | |
| end | |
| if !isa(α, Const) && !isa(C, Const) | |
| if iszero(β.val) && !iszero(α.val) | |
| # this is probably quite a common case, so maybe worth specializing | |
| mul!(C.val, A.val, B.val, α.val, β.val) | |
| add!(C.dval, C.val, α.dval / α.val) | |
| else | |
| AB = A.val * B.val | |
| add!(C.val, AB, α.val, β.val) | |
| add!(C.dval, AB, α.dval) | |
| end | |
| else | |
| mul!(C.val, A.val, B.val, α.val, β.val) | |
| end |
There was a problem hiding this comment.
Something went wrong, this is a suggestion for the whole function body above.
There was a problem hiding this comment.
Also, I don't think we need project_mul! here, since C should be complex as soon as any of A, B, α or β is complex.
Can't you pass in real-valued C and complex valued other arguments if you are using the direct in-place method?
There was a problem hiding this comment.
Tried the suggestion locally and it worked, thanks!
There was a problem hiding this comment.
There's no forcing, but the function will throw InexactError because setindex! on C will have to convert for any non-real value.
I guess you could in principle come up with an evil case where the eltype of one of the inputs is complex but the values are purely real, which then might work and give issues here?
There was a problem hiding this comment.
Yeah exactly, suppose I generate C, A, B all real then get Ac = complex(A) which is of course all "real" valued. But this doesn't need to be solved in the forward rules, this PR has been expanded enough
|
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 |
3 participants
Trying to make these a little more manageable and pick up the fwd rules where possible