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, …

In fact, I am trying to apply it to automatic differentiation for finite element analysis (FEA). So when I use it, it will be an appropriate sparse matrix. However, what I noticed is that the linear solver part takes too much time, even at small sizes. That is why I present this question.

There has been some discussion about the impact of applying automatic differentiation on Julia’s coding, so here are my thoughts. I have created the same project with the following three types of projects. The project involves numerical simulations using the finite element method (FEM). Enzyme.j…

Does Sparspak.jl + ForwardDiff work for you ?

[image] ToPo: I still think the solution is to write your own custom rrules. Would replacing them with custom rrules also reduce memory? Yes, in my experience often a custom rrule can be both faster and more memory-efficient than reverse-mode AD, assuming you know what you are doing in optim…

Did you try LinearSolve.jl? It adds rules for Zygote and Enzyme and wraps pretty much every linear solver that people use in practice. That should handle this just fine?

[image] ChrisRackauckas: It adds rules for Zygote and Enzyme and wraps pretty much every linear solver that people use in practice. That should handle this just fine? To compute \nabla_x \Vert A(x)^{-1} b \Vert with current AD tools where A(x) is sparse, you still have the problem of trying t…

Thank you. I understand. I will try to define my own rrule for the areas where I want to improve efficiency.

[image] j-fu: Does Sparspak.jl + ForwardDiff work for you ? I tried small tests with Sparspak.jl + ForwardDiff.jl + ExtendableSparse.jl in the last week, but it didin’t work. Maybe additional testing will be done to show the new topic. Would it work well without Exrendable?

It should. Basically it should work as well with ExtendableSparse. Could you open an issue there (possibly with an MWE ?)

[image] stevengj: To compute \nabla_x \Vert A(x)^{-1} b \Vert∇x∥A(x)−1b∥\nabla_x \Vert A(x)^{-1} b \Vert with current AD tools where A(x)A(x)A(x) is sparse, you still have the problem of trying to differentiate through the sparse-matrix constructor , no? I checked my tests and yes sparse is m…

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

Specific Domains Numerics

stevengj March 20, 2024, 3:40pm 15

To compute \nabla_x \Vert A(x)^{-1} b \Vert with current AD tools where A(x) is sparse, you still have the problem of trying to differentiate through the sparse-matrix constructor, no?

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

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

Related topics