I have confirmed that the standard Julia solver (backslash) can be used in the automatic differentiation of Zygote. Below is a sample program. (This program calculates the gradient.) However, while the forward computation of the evaluation function is fast, the computation of the automatic derivat…

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. using LinearAlgebra, LinearSolve, …

[image] wsmoses: A previous pain point we’ve since mostly resolved, is now we have fast versions of most common BLAS functions (gemv, gemm, dot, etc), however we haven’t done this for LAPack yet (which Julia will call from det/etc). I’d expand it to say, the C boundaries are what’s largely no…

[image] ChrisRackauckas: I checked my tests and yes sparse is missing an adjoint. As I explained in my linked post, it’s not simply a matter of a missing rule — if you don’t construct an implicit representation of the tangent vector upstream, the input to your adjoint/rrule/pullback/vJp will …

[image] stevengj: But do any current AD systems do this by default without manual intervention? It requires manual intervention but maybe ImplicitDifferentiation.jl can help here. It seems constructing the sparse matrix A(x) makes sense for solving A(x) \ b optimally. However, the optimality …

Note I just checked, this all is a non-issue with Enzyme as it differentiates fine: using Enzyme, SparseArrays function f1(A) sum(SparseMatrixCSC(A)) end A = rand(2,2) dA = zeros(2,2) Enzyme.autodiff(Reverse, f1, Duplicated(A, dA)) @show dA 2×2 Matrix{Float64}: 1.0 1.0 1.0 1.0 I = [1, 4, …

[image] ChrisRackauckas: Note I just checked, this all is a non-issue with Enzyme as it differentiates fine: How do you know it’s not constructing a dense matrix as an intermediate step? [image] gdalle: It requires manual intervention but maybe ImplicitDifferentiation.jl can help here. …

[image] stevengj: It’s not that this can’t be done — I pointed out the required implicit representation in one common case in my post, and @mohamed82008 gave some example code. For reference, Zygote.jl: How to get the gradient of sparse matrix - #12 by mohamed82008

[image] stevengj: How do you know it’s not constructing a dense matrix as an intermediate step? Actually, it looks good — I tried it with a {10}^6 \times {10}^6 sparse matrix, which would have run out of memory if it tried to allocate something dense, and it seems to work fine for my simple e…

I think a related issue is also what people mean when they want a gradient wrt to a sparse matrix. Do I mean the gradient wrt the whole matrix including its structural zeros which is the mathematically natural way to think about this? Or do I mean the gradient wrt to the matrix’s structural non-zero…

[image] stevengj: As I explained in my linked post, the same issue potentially arises with something as simple as f(x) = u^T A(x) vf(x)=uTA(x)vf(x) = u^T A(x) v. I know, but the point is to leverage A being a sparse matrix when we do A(x) \ b, because we have efficient routines for that. So t…

[image] stevengj: how does Enzyme represent the tangent vector for input into the sparse adjoint? Related discussion: Supporting covariant derivatives opened 12:25PM - 07 Mar 24 UTC (UTC) [MasonProt…

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

Specific Domains Numerics

mohamed82008 March 21, 2024, 1:37pm 26

For reference, Zygote.jl: How to get the gradient of sparse matrix - #12 by mohamed82008

Topic		Replies	Views
How to efficiently differentiate backslash operator for sparse matrix? Numerics differentiation , forwarddiff , sparse , juafem , finite-element	23	2824	May 1, 2023
Zygote.jl: How to get the gradient of sparse matrix General Usage question , package , differentiation , zygote , reversediff	11	1424	June 13, 2023
Automatic differentiation, adjoint methods, and linear systems of equations General Usage	4	1235	May 8, 2019
Sparse jacobians of matrix models Machine Learning question , documentation , performance , optimization , sparse	37	2110	July 3, 2022
Zygote much slower than JAX for automatic differentiation of energy Machine Learning performance , flux , zygote , jax , lux	22	1737	May 15, 2024

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

Related topics