giggleliu / backwardslinalg.jl

Auto differentiation over linear algebras (a Zygote extension)

Hi there 👋

backwardslinalg.jl's Issues

Working with QR gradients

Hi @GiggleLiu,

I was wondering about this use case

using Zygote
using BackwardsLinalg
gradient(x -> sum(qr(x).Q), rand(3,3)) 

but this gives an error

julia> gradient(x -> qr(x).Q |> sum, rand(3,3))
ERROR: MethodError: no method matching getindex(::Tuple{Array{Float64,2},Array{Float64,2}}, ::Symbol)
Closest candidates are:
  getindex(::Tuple, ::Int64) at tuple.jl:24
  getindex(::Tuple, ::Real) at tuple.jl:25
  getindex(::Tuple, ::AbstractUnitRange{#s66} where #s66<:Real) at range.jl:288
  ...
Stacktrace:
 [1] adjoint at /home/meach/.julia/packages/Zygote/fuj2C/src/lib/lib.jl:81 [inlined]
 [2] _forward at /home/meach/.julia/packages/Zygote/fuj2C/src/lib/grad.jl:44 [inlined]
 [3] _forward(::Zygote.Context, ::typeof(Zygote.literal_getproperty), ::Tuple{Array{Float64,2},Array{Float64,2}}, ::Val{:Q}) at /home/meach/.julia/packages/Zygote/fuj2C/src/lib/lib.jl:174
 [4] #19 at ./REPL[14]:1 [inlined]
 [5] _forward(::Zygote.Context, ::getfield(Main, Symbol("##19#20")), ::Array{Float64,2}) at /home/meach/.julia/packages/Zygote/fuj2C/src/compiler/interface2.jl:0
 [6] _forward(::Function, ::Array{Float64,2}) at /home/meach/.julia/packages/Zygote/fuj2C/src/compiler/interface.jl:31
 [7] forward(::Function, ::Array{Float64,2}) at /home/meach/.julia/packages/Zygote/fuj2C/src/compiler/interface.jl:37
 [8] gradient(::Function, ::Array{Float64,2}) at /home/meach/.julia/packages/Zygote/fuj2C/src/compiler/interface.jl:46
 [9] top-level scope at REPL[14]:1

Is there any way to fix this?

Consider string literal interface

U, S, V = svd"ijkl->ijs,s,skl"(A; tol=?, D =?)

This would be helpful in tensor network algorithms.

QR with pivoting?

Does this support qr! with pivoting for complex, non-square matricies?

using BackwardsLinalg
f(x) = qr!(x, val(true)).Q
A = rand(4,8) .+ 1.5*im
f'(A)

For registration

register

@JuliaRegistrator register

GPU support

return `SVD` and `QR` types to stay compatible with julialang

Merge into ChainRules?

It's great to see someone taking AD + linear algebra seriously. @GiggleLiu would you be open to merging this stuff into ChainRules at some point in the not too distant future? It would be a great help!

References

Paper

• symeig, lq, et. al
  https://arxiv.org/abs/1710.08717
• svd
  https://j-towns.github.io/papers/svd-derivative.pdf
• rsvd
  https://epubs.siam.org/doi/10.1137/090771806
• Einsum
  https://github.com/pytorch/pytorch/blob/4404762d7dd955383acee92e6f06b48144a0742e/aten/src/ATen/native/Linear.cpp#L125
• Lei-Wang's autodiff in TN paper (about increasing numerical stability)
  To be announced

Tensorflow

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/ops/linalg_grad.py

Julia linalg autodiff

https://github.com/invenia/Nabla.jl

register the package

You need to add me as collaborator, so I will have the access.

Batched SVD

lapack exception (2) backward qr diff

While random matrices work, I get lapack exceptions in trtrs when using it in my application.
fpeps.zip

`using Zygote,BackwardsLinalg,fpeps

width = 4;height = 4;
D = 2;d = 2;chi = 5;

peps = Array{Array{ComplexF64,5},2}(undef,width,height);
for i = 1:width
for j = 1:height
peps[i,j] = rand(ComplexF64,D,D,D,D,d);
end
end

boundaries = fpeps.gen_boundaries(peps,chi);
n1 = boundaries[1][1]; n2 = boundaries[1][2];

fun(a) = real(sum(sum.(fpeps.north_vomps(n1,peps[1,:],a))))
fun'(n2)`

Is there anything I should know about numerical stability problems (assuming the lapack error means that trtrs didn't converge). How can I go about debugging this further? The matrix passed to trtrs did became spectacularly quite ill-conditioned in my example... I can try to reduce this further to a smaller minimal example?

If I set the parameters to make sure that we are effectively working with 1x1 matrices then I get errors later on with zygote (I'm assuming it's unrelated to this package but I'm not sure as the stacktrace effectively says nothing).