Make the rrule for 3-arg dot lazy

[mohamed82008]

Currently this line

ChainRules.jl/src/rulesets/LinearAlgebra/dense.jl

Line 38 in 9f1817a

dA = @thunk project_A(ΔΩ .* x .* adjoint(y))

leads to a loss of structure because it constructs a dense matrix from 2 vectors. We should make this a lazy outer product (like in SciML/LinearSolve.jl#484) such that the m x n matrix's co-tangent only requires O(m + n) memory. See https://discourse.julialang.org/t/zygote-jl-how-to-get-the-gradient-of-sparse-matrix/59067/12 for the potential gain and suggested implementation. If the idea gets a green light, I will open a PR.

[oxinabox]

I guess the key question is are we willing to add LazyArrays as a dependency?

[gdalle]

Related: https://discourse.julialang.org/t/zygote-much-slower-than-jax-for-automatic-differentiation-of-energy/114239/5

[oxinabox]

Since this seems to have come up twice in two different contexts i think we can say that this is worth it and we should add LazyArrays as a dependency.
@mohamed82008 are you able to make the PR?

[mohamed82008]

I am able to yes. I will do it in the weekend unless someone wants to beat me to it. It's basically exactly this implementation https://discourse.julialang.org/t/zygote-jl-how-to-get-the-gradient-of-sparse-matrix/59067/12 plus unit tests.

[albertomercurio]

Hello, are there any news about this? I would like to implement reverse diff to dot(x, A, x) function with A sparse, but at the moment it converts the A matrix into a dense one.

[mohamed82008]

#796