Fix Enzyme integration following minimal example about linearsolve.jl

[ChrisRackauckas]

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 linsolve object.

@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.

[wsmoses]

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: ***@***.***>