Using Enzyme with LinearSolve, scary errors!

General Usage
1

I have a MWE of using Enzyme with linearsolve that fails, giving first a warning about using a BLAS fallback (which gets shoved out of my terminal by what follows), then a huge amount of Enzyme TypeAnalysisDepthLimit warnings, and the stacktrace below.

It is obvious that work has been done to make Enzyme and LinearSolve compatible, due to the extension that exists, but I haven’t found any examples of it in use. I would really appreciate being pointed to one. Bonus points if it can be done non-allocating!

using Enzyme
using LinearSolve

function testls(A, b, u)
    oa = OperatorAssumptions(true, condition = LinearSolve.OperatorCondition.WellConditioned)
    prob = LinearProblem{true}(A, b, u0=u)
    linsolve = init(prob, oa)
    solve(linsolve)
    return nothing
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))


ERROR: Enzyme execution failed.
Enzyme: Augmented forward pass custom rule Tuple{EnzymeCore.EnzymeRules.ConfigWidth{1, true, false, (false, true)}, Const{typeof(solve!)}, Type{Const{SciMLBase.LinearSolution{Float64, 1, Vector{Float64}, Nothing, LinearSolve.DefaultLinearSolver, LinearSolve.LinearCache{Matrix{Float64}, Vector{Float64}, Vector{Float64}, SciMLBase.NullParameters, LinearSolve.DefaultLinearSolver, LinearSolve.DefaultLinearSolverInit{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, LinearAlgebra.QRCompactWY{Float64, Matrix{Float64}, Matrix{Float64}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Vector{Int64}}, Nothing, Nothing, Nothing, LinearAlgebra.SVD{Float64, Float64, Matrix{Float64}, Vector{Float64}}, LinearAlgebra.Cholesky{Float64, Matrix{Float64}}, LinearAlgebra.Cholesky{Float64, Matrix{Float64}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int32}}, Base.RefValue{Int32}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Base.RefValue{Int64}}, LinearAlgebra.QRPivoted{Float64, Matrix{Float64}, Vector{Float64}, Vector{Int64}}, Krylov.CraigmrSolver{Float64, Float64, Vector{Float64}}, Krylov.LsmrSolver{Float64, Float64, Vector{Float64}}}, IdentityOperator, IdentityOperator, Float64, Bool}, Nothing}}}, Duplicated{LinearSolve.LinearCache{Matrix{Float64}, Vector{Float64}, Vector{Float64}, SciMLBase.NullParameters, LinearSolve.DefaultLinearSolver, LinearSolve.DefaultLinearSolverInit{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, LinearAlgebra.QRCompactWY{Float64, Matrix{Float64}, Matrix{Float64}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Vector{Int64}}, Nothing, Nothing, Nothing, LinearAlgebra.SVD{Float64, Float64, Matrix{Float64}, Vector{Float64}}, LinearAlgebra.Cholesky{Float64, Matrix{Float64}}, LinearAlgebra.Cholesky{Float64, Matrix{Float64}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int32}}, Base.RefValue{Int32}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Base.RefValue{Int64}}, LinearAlgebra.QRPivoted{Float64, Matrix{Float64}, Vector{Float64}, Vector{Int64}}, Krylov.CraigmrSolver{Float64, Float64, Vector{Float64}}, Krylov.LsmrSolver{Float64, Float64, Vector{Float64}}}, IdentityOperator, IdentityOperator, Float64, Bool}}} return type mismatch, expected EnzymeCore.EnzymeRules.AugmentedReturn{SciMLBase.LinearSolution{Float64, 1, Vector{Float64}, Nothing, LinearSolve.DefaultLinearSolver, LinearSolve.LinearCache{Matrix{Float64}, Vector{Float64}, Vector{Float64}, SciMLBase.NullParameters, LinearSolve.DefaultLinearSolver, LinearSolve.DefaultLinearSolverInit{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, LinearAlgebra.QRCompactWY{Float64, Matrix{Float64}, Matrix{Float64}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Vector{Int64}}, Nothing, Nothing, Nothing, LinearAlgebra.SVD{Float64, Float64, Matrix{Float64}, Vector{Float64}}, LinearAlgebra.Cholesky{Float64, Matrix{Float64}}, LinearAlgebra.Cholesky{Float64, Matrix{Float64}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int32}}, Base.RefValue{Int32}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Base.RefValue{Int64}}, LinearAlgebra.QRPivoted{Float64, Matrix{Float64}, Vector{Float64}, Vector{Int64}}, Krylov.CraigmrSolver{Float64, Float64, Vector{Float64}}, Krylov.LsmrSolver{Float64, Float64, Vector{Float64}}}, IdentityOperator, IdentityOperator, Float64, Bool}, Nothing}, Nothing, Any} found EnzymeCore.EnzymeRules.AugmentedReturn{SciMLBase.LinearSolution{Float64, 1, Vector{Float64}, Nothing, LinearSolve.DefaultLinearSolver, LinearSolve.LinearCache{Matrix{Float64}, Vector{Float64}, Vector{Float64}, SciMLBase.NullParameters, LinearSolve.DefaultLinearSolver, LinearSolve.DefaultLinearSolverInit{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, LinearAlgebra.QRCompactWY{Float64, Matrix{Float64}, Matrix{Float64}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Vector{Int64}}, Nothing, Nothing, Nothing, LinearAlgebra.SVD{Float64, Float64, Matrix{Float64}, Vector{Float64}}, LinearAlgebra.Cholesky{Float64, Matrix{Float64}}, LinearAlgebra.Cholesky{Float64, Matrix{Float64}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int32}}, Base.RefValue{Int32}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Base.RefValue{Int64}}, LinearAlgebra.QRPivoted{Float64, Matrix{Float64}, Vector{Float64}, Vector{Int64}}, Krylov.CraigmrSolver{Float64, Float64, Vector{Float64}}, Krylov.LsmrSolver{Float64, Float64, Vector{Float64}}}, IdentityOperator, IdentityOperator, Float64, Bool}, Nothing}, SciMLBase.LinearSolution{Float64, 1, Vector{Float64}, Nothing, LinearSolve.DefaultLinearSolver, LinearSolve.LinearCache{Matrix{Float64}, Vector{Float64}, Vector{Float64}, SciMLBase.NullParameters, LinearSolve.DefaultLinearSolver, LinearSolve.DefaultLinearSolverInit{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, LinearAlgebra.QRCompactWY{Float64, Matrix{Float64}, Matrix{Float64}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Vector{Int64}}, Nothing, Nothing, Nothing, LinearAlgebra.SVD{Float64, Float64, Matrix{Float64}, Vector{Float64}}, LinearAlgebra.Cholesky{Float64, Matrix{Float64}}, LinearAlgebra.Cholesky{Float64, Matrix{Float64}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int32}}, Base.RefValue{Int32}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Base.RefValue{Int64}}, LinearAlgebra.QRPivoted{Float64, Matrix{Float64}, Vector{Float64}, Vector{Int64}}, Krylov.CraigmrSolver{Float64, Float64, Vector{Float64}}, Krylov.LsmrSolver{Float64, Float64, Vector{Float64}}}, IdentityOperator, IdentityOperator, Float64, Bool}, Nothing}, Tuple{Vector{Float64}, Vector{Float64}, LinearSolve.LinearCache{Matrix{Float64}, Vector{Float64}, Vector{Float64}, SciMLBase.NullParameters, LinearSolve.DefaultLinearSolver, LinearSolve.DefaultLinearSolverInit{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, LinearAlgebra.QRCompactWY{Float64, Matrix{Float64}, Matrix{Float64}}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Vector{Int64}}, Nothing, Nothing, Nothing, LinearAlgebra.SVD{Float64, Float64, Matrix{Float64}, Vector{Float64}}, LinearAlgebra.Cholesky{Float64, Matrix{Float64}}, LinearAlgebra.Cholesky{Float64, Matrix{Float64}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int32}}, Base.RefValue{Int32}}, Tuple{LinearAlgebra.LU{Float64, Matrix{Float64}, Vector{Int64}}, Base.RefValue{Int64}}, LinearAlgebra.QRPivoted{Float64, Matrix{Float64}, Vector{Float64}, Vector{Int64}}, Krylov.CraigmrSolver{Float64, Float64, Vector{Float64}}, Krylov.LsmrSolver{Float64, Float64, Vector{Float64}}}, IdentityOperator, IdentityOperator, Float64, Bool}, Tuple{Matrix{Float64}}, Tuple{Vector{Float64}}}}
Stacktrace:
 [1] #solve#96
   @ ./deprecated.jl:105

Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/jOGYG/src/compiler.jl:1317
  [2] #solve#96
    @ ./deprecated.jl:105
  [3] testls
    @ ~/Dropbox (Cambridge University)/shared_Daniel/code/RobotDynamics.jl/examples/gradients/enzyme_experiments.jl:178 [inlined]
  [4] diffejulia_testls_4370wrap
    @ ~/Dropbox (Cambridge University)/shared_Daniel/code/RobotDynamics.jl/examples/gradients/enzyme_experiments.jl:0
  [5] macro expansion
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/jOGYG/src/compiler.jl:5306 [inlined]
  [6] enzyme_call(::Val{…}, ::Ptr{…}, ::Type{…}, ::Type{…}, ::Val{…}, ::Type{…}, ::Type{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Duplicated{…}, ::Duplicated{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/jOGYG/src/compiler.jl:4984
  [7] (::Enzyme.Compiler.CombinedAdjointThunk{…})(::Const{…}, ::Duplicated{…}, ::Vararg{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/jOGYG/src/compiler.jl:4926
  [8] autodiff(::ReverseMode{…}, ::Const{…}, ::Type{…}, ::Duplicated{…}, ::Vararg{…})
    @ Enzyme ~/.julia/packages/Enzyme/jOGYG/src/Enzyme.jl:215
  [9] autodiff(::ReverseMode{…}, ::Const{…}, ::Duplicated{…}, ::Duplicated{…}, ::Vararg{…})
    @ Enzyme ~/.julia/packages/Enzyme/jOGYG/src/Enzyme.jl:238
 [10] autodiff(::ReverseMode{…}, ::typeof(testls), ::Duplicated{…}, ::Duplicated{…}, ::Vararg{…})
    @ Enzyme ~/.julia/packages/Enzyme/jOGYG/src/Enzyme.jl:224
 [11] top-level scope
    @ ~/Dropbox (Cambridge University)/shared_Daniel/code/RobotDynamics.jl/examples/gradients/enzyme_experiments.jl:191
Some type information was truncated. Use `show(err)` to see complete types.
2

That looks to be an error in the LinearSolve.jl EnzymeRule.

Open an issue on LinearSolve.jl and perhaps consider using the builtin solvers in the interim?

3

Thanks for the reply - I’ve had a go with the builtin solvers and cooked up a minimal example to show what’s going on.

The first case seems to work, but the other 2 fail. this should be fine for now, but can you explain to me why this is the case? Is there an enzyme rule written somewhere for \ but not for the others?

function testsolve_backslash!(A, b, u)
    copyto!(u, A\b)
    return u[1]
end

function testsolve_lu!(A, b, u)
    factorization = lu!(A) # This line allocates...
    ldiv!(factorization, copyto!(u, b))
    return u[1]
end

function testsolve_cholesky!(A, b, u)
    factorization = cholesky!(A) # This is non allocating
    ldiv!(factorization, copyto!(u, b))
    return u[1]
end


A = [3. 1.; 1. 2.]
b = [1., 2.]
u = zero(b)

for method in (testsolve_backslash!, testsolve_lu!, testsolve_cholesky!)
    _A, _b, _u = copy(A), copy(b), copy(u)
    @time method(_A, _b, _u)
    @assert _u ≈ A\b
end

A_enz, b_enz, u_enz = Duplicated(copy(A), zero(A)), Duplicated(copy(b), zero(b)), Duplicated(copy(u), zero(u))
# This works
Enzyme.autodiff(Reverse, testsolve_backslash!, A_enz, b_enz, u_enz)
# This errors:
Enzyme.autodiff(Reverse, testsolve_lu!, A_enz, b_enz, u_enz)
# This errors:
Enzyme.autodiff(Reverse, testsolve_cholesky!, A_enz, b_enz, u_enz)
4

Yeah the first has a rule, but I don’t think the later to do.

@michel_schanen and @simsurace have been looking at linear solve rules, tagging them for visibility. Maybe file an issue for what doesn’t work right now?

5

I don’t know what that error message is supposed to mean though. Is it that the return type is somehow different?

That shouldn’t matter because it doesn’t actually differentiate the solve, it boxes that part.

6

Removing the giant sciml types to make it easier to read

Enzyme: Augmented forward pass custom rule [input types for the function], return type mismatch, expected EnzymeCore.EnzymeRules.AugmentedReturn{..., Nothing, Any} found EnzymeCore.EnzymeRules.AugmentedReturn{..., ..., ...}

Looks like the rule returned a shadow return, where it was not requested. You can do EnzymeCore.needs_shadow(config) to determine whether you should return a shadow. Enzyme.jl/lib/EnzymeCore/src/rules.jl at e665d895b7eb0cdb7c3bb230d9272425a1918219 · EnzymeAD/Enzyme.jl · GitHub

7

Can someone open an issue for this? I don’t want to forget/lose it

8

Does the third example work with Fix rules for `cholesky` and `ldiv!` on `Cholesky` by simsurace · Pull Request #1307 · EnzymeAD/Enzyme.jl · GitHub?

9

Indeed it doesn’t, because the rule is only for cholesky, but not cholesky!. Apart from that, there are still errors in that PR, need to revisit once more.

There is currently no rule for lu!.

10

I’ve made an issue here - it’s my first every on an open project, let me know if I’ve messed up anything.

11

That looks good, thanks

1 Like