# How do you speed up the linear sparse solver in Zygote?

**URL:** <https://discourse.julialang.org/t/how-do-you-speed-up-the-linear-sparse-solver-in-zygote/111801>\
**Category:** Numerics\
**Tags:** question, zygote, linearsolve, sparsearrays\
**Created:** [March 19, 2024, 7:15am UTC](https://discourse.julialang.org/t/how-do-you-speed-up-the-linear-sparse-solver-in-zygote/111801 "2024-03-19T07:15:50Z")\
**Posts on this page:** 5\
**Page:** 3

<div class="post-metadata">

**Author:** ![mohamed82008](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mohamed82008/32/18171_2.png) [@mohamed82008](https://discourse.julialang.org/u/mohamed82008)\
**Post date:** [March 22, 2024, 4:30pm UTC](https://discourse.julialang.org/t/how-do-you-speed-up-the-linear-sparse-solver-in-zygote/111801/41 "2024-03-22T16:30:45Z")

</div>

Thinking some more about the example I posted above. In practice, one wouldn’t write `inv` in the code at all. And one can define a rule for the linear solve directly so u and I should be a single function call. This means that the following line can be optimized in the `rrule` of the linear solve (reusing the factorization, reusing the linear solve solution and doing the lazy 1 rank matrix representation).

dK=−(K^{−1} \cdot du) \cdot (K^{-1} \cdot f)^T

In fact, maybe all that’s needed is to make this line [LinearSolve.jl/src/adjoint.jl at main · SciML/LinearSolve.jl · GitHub](https://github.com/SciML/LinearSolve.jl/blob/main/src/adjoint.jl#L79) a lazy multiplication returning a 1-rank matrix instead of a dense one. @ChrisRackauckas would you be open to a PR for this?

Then one can define a rule for dxi=tr(dK^T \cdot K\_i) or `dot(A, B)` in Julia where `A` is a 1-rank matrix and B is a `SparseMatrixCSC`. Then we can get near optimal performance in the above case using Zygote. Not sure if Enzyme rules support lazy arrays for the adjoint. This may be easier than I thought initially without too many changes.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [March 22, 2024, 9:49pm UTC](https://discourse.julialang.org/t/how-do-you-speed-up-the-linear-sparse-solver-in-zygote/111801/42 "2024-03-22T21:49:22Z")

</div>

> [@mohamed82008](#):
>
> In fact, maybe all that’s needed is to make this line [LinearSolve.jl/src/adjoint.jl at main · SciML/LinearSolve.jl · GitHub](https://github.com/SciML/LinearSolve.jl/blob/main/src/adjoint.jl#L79) a lazy multiplication returning a 1-rank matrix instead of a dense one. @ChrisRackauckas would you be open to a PR for this?

I’d be open to a PR for this. I don’t know what would happen if it’s not matching the type of A but it’s worth investigating.

---

<div class="post-metadata">

**Author:** ![mohamed82008](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mohamed82008/32/18171_2.png) [@mohamed82008](https://discourse.julialang.org/u/mohamed82008)\
**Post date:** [March 23, 2024, 10:25am UTC](https://discourse.julialang.org/t/how-do-you-speed-up-the-linear-sparse-solver-in-zygote/111801/43 "2024-03-23T10:25:14Z")

</div>

> <https://github.com/SciML/LinearSolve.jl/pull/484>
>
> \## Checklist
> 
> \- \[x\] Appropriate tests were added
> \- \[\] Any code changes were …done in a way that does not break public API
> \- \[\] All documentation related to code changes were updated
> \- \[x\] The new code follows the
> \[contributor guidelines\](https://github.com/SciML/.github/blob/master/CONTRIBUTING.md), in particular the \[SciML Style Guide\](https://github.com/SciML/SciMLStyle) and
> \[COLPRAC\](https://github.com/SciML/COLPRAC).
> \- \[\] Any new documentation only uses public API
>   
> \## Additional context
> 
> This PR implements the suggestion in https://discourse.julialang.org/t/how-do-you-speed-up-the-linear-sparse-solver-in-zygote/111801/41?u=mohamed82008. To be safe, I also bumped the major version because the output type of \`Zygote.gradient\` wrt the matrix is changed in this PR. No new documentation was added as this PR is just a performance improvement.

---

<div class="post-metadata">

**Author:** ![ToPo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/topo/32/205868_2.png) [@ToPo](https://discourse.julialang.org/u/ToPo)\
**Post date:** [March 23, 2024, 3:33pm UTC](https://discourse.julialang.org/t/how-do-you-speed-up-the-linear-sparse-solver-in-zygote/111801/44 "2024-03-23T15:33:19Z")

</div>

Thank you! I will refer to it.

---

<div class="post-metadata">

**Author:** ![mohamed82008](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mohamed82008/32/18171_2.png) [@mohamed82008](https://discourse.julialang.org/u/mohamed82008)\
**Post date:** [March 24, 2024, 7:18am UTC](https://discourse.julialang.org/t/how-do-you-speed-up-the-linear-sparse-solver-in-zygote/111801/45 "2024-03-24T07:18:47Z")

</div>

The PR above just got merged which means that in the **next version** of LinearSolve.jl if you use Zygote.jl (or any ChainRules.jl based AD package) to get the gradient of `dot(a, A \ b)` wrt the matrix `A`, you will get a lazy outer product. This is true when `A` is dense.

```julia
using LinearAlgebra, LinearSolve, Zygote

function invquad(a, A, b)
    prob = LinearProblem(A, b)
    sol = solve(
        prob,
        LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.RFLUFactorization),
    )
    return dot(a, sol.u)
end

n = 100; A = rand(n, n); b1 = rand(n); b2 = rand(n);

db1, dA, db2 = Zygote.gradient(invquad, b1, A, b2);

Base.summarysize(dA)
# 1752

Base.summarysize(A)
# 80040

```

If `A` is sparse, `Zygote.pullback` gives you the correct lazy outer product but `Zygote.gradient` projects that to be a sparse matrix with the same sparsity structure as the input which is arguably incorrect ([Gradient wrt to a sparse matrix is mathematically wrong · Issue #1507 · FluxML/Zygote.jl · GitHub](https://github.com/FluxML/Zygote.jl/issues/1507)) because the gradient is mathematically defined for the structural zeros. The fact that they are structural zeros is an implementation detail. The projection imo should be done by the user after the `gradient` call if needed. Anyways, if you use `pullback` directly, you are safe.

```julia
function sinvquad(a, A, b)
    prob = LinearProblem(A, b)
    sol = solve(prob)
    return dot(a, sol.u)
end

sA = sparse(A)
dsA = Zygote.pullback(sinvquad, b1, sA, b2)[2](1.0)[2]

Base.summarysize(dsA)
# 1752

Base.summarysize(sA)
# 160968

```

I think the above PR was the first step towards preserving structure in rules and using lazy representations where possible. The next step is [Make the rrule for 3-arg dot lazy · Issue #788 · JuliaDiff/ChainRules.jl · GitHub](https://github.com/JuliaDiff/ChainRules.jl/issues/788).

[Previous page](https://discourse.julialang.org/t/how-do-you-speed-up-the-linear-sparse-solver-in-zygote/111801.md?page=2)