Fix Enzyme integration following minimal example #479
Labels
bug
Something isn't working
Comments
|
Using the return works: using Enzyme
using LinearSolve
function testls(A, b, u)
oa = OperatorAssumptions(true, condition = LinearSolve.OperatorCondition.WellConditioned)
prob = LinearProblem(A, b)
linsolve = init(prob, LUFactorization(), assumptions = oa)
cache =solve!(linsolve)
sum(cache.u)
end
A = [1. 2.; 3. 4.]
b = [1., 2.]
u = zero(b)
dA = deepcopy(A)
db = deepcopy(b)
du = deepcopy(u)
@time testls(A, b, u)
Enzyme.autodiff(Reverse, testls, Duplicated(A, dA), Duplicated(b, db), Duplicated(u, du))it's just when I change it to use the mutation directly: function testls(A, b, u)
oa = OperatorAssumptions(true, condition = LinearSolve.OperatorCondition.WellConditioned)
prob = LinearProblem(A, b)
linsolve = init(prob, LUFactorization(), assumptions = oa)
solve!(linsolve)
sum(linsolve.u)
end
Enzyme.autodiff(Reverse, testls, Duplicated(A, dA), Duplicated(b, db), Duplicated(u, du))that it then fails. Which is weird because the return is just the mutated @wsmoses is this an Enzyme limitation that if the return is ignored it needs to assume |
|
No this isn’t an Enzyme limitation. Because the return isn’t used Enzyme
specifies in the config to the rule that you should not return the result /
shadow (accessible via needsPrimal/needsshadow). The rule does in that
case, which should be a simple fix to check that and conditionally return
To be clear we should extend Enzyme to still work even in this case where
more data is returned than requested, but the above fix should work and
would be preferred.
…On Tue, Mar 12, 2024 at 6:12 AM Christopher Rackauckas < ***@***.***> wrote:
Using the return works:
using Enzymeusing LinearSolve
function testls(A, b, u)
oa = OperatorAssumptions(true, condition = LinearSolve.OperatorCondition.WellConditioned)
prob = LinearProblem(A, b)
linsolve = init(prob, LUFactorization(), assumptions = oa)
cache =solve!(linsolve)
sum(cache.u)end
A = [1. 2.; 3. 4.]
b = [1., 2.]
u = zero(b)
dA = deepcopy(A)
db = deepcopy(b)
du = ***@***.*** testls(A, b, u)
Enzyme.autodiff(Reverse, testls, Duplicated(A, dA), Duplicated(b, db), Duplicated(u, du))
it's just when I change it to use the mutation directly:
function testls(A, b, u)
oa = OperatorAssumptions(true, condition = LinearSolve.OperatorCondition.WellConditioned)
prob = LinearProblem(A, b)
linsolve = init(prob, LUFactorization(), assumptions = oa)
solve!(linsolve)
sum(linsolve.u)end
Enzyme.autodiff(Reverse, testls, Duplicated(A, dA), Duplicated(b, db), Duplicated(u, du))
that it then fails. Which is weird because the return is just the mutated
linsolve object.
@wsmoses <https://github.com/wsmoses> is this an Enzyme limitation that
if the return is ignored it needs to assume nothing return and thus has a
type mismatch? What we can do with LinearSolve.jl is just make solve!
return nothing since it's mutating (this would be breaking, but might be
a good change) and that would help make sure all codes support Enzyme, or
if there's a way to handle this case that would be good to know.
—
Reply to this email directly, view it on GitHub
<#479 (comment)>,
or unsubscribe
<https://github.com/notifications/unsubscribe-auth/AAJTUXD5OVKS56PUG2CGQGLYX35MBAVCNFSM6AAAAABEFGU566VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSOJRGYZDCNRUGA>
.
You are receiving this because you were mentioned.Message ID:
***@***.***>
|
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Enzyme autodiff fails to work with LinearSolve.
Minimal Reproducible Example 👇
Additional context
See discourse thread
The text was updated successfully, but these errors were encountered: