# Autodiff : Custom derivatives of vector functions for linear algebra

**URL:** <https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826>\
**Category:** Specific Domains\
**Tags:** linearalgebra\
**Created:** [April 6, 2019, 3:27am UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826 "2019-04-06T03:27:26Z")\
**Posts on this page:** 16\
**Page:** 1

<div class="post-metadata">

**Author:** ![sdewaele](https://avatars.discourse-cdn.com/v4/letter/s/7ab992/32.png) [@sdewaele](https://discourse.julialang.org/u/sdewaele)\
**Post date:** [April 6, 2019, 3:27am UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/1 "2019-04-06T03:27:26Z")

</div>

Is there a solution for automatic differentiation through code that contains a number of linear algebra operations including, for example, [lyap](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/index.html#LinearAlgebra.lyap) and the matrix exponential exponential ([exp](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/index.html#Base.exp-Tuple%7BUnion%7BDenseArray%7B#s36,2%7D,%20ReinterpretArray%7B#s36,2,S,A%7D%20where%20S%20where%20A%3C:Union%7BSubArray%7BT,N,A,I,true%7D%20where%20I%3C:Union%7BTuple%7BVararg%7BReal,N%7D%20where%20N%7D,%20Tuple%7BAbstractUnitRange,Vararg%7BAny,N%7D%20where%20N%7D%7D%20where%20A%3C:DenseArray%20where%20N%20where%20T,%20DenseArray%7D,%20ReshapedArray%7B#s36,2,A,MI%7D%20where%20MI%3C:Tuple%7BVararg%7BSignedMultiplicativeInverse%7BInt64%7D,N%7D%20where%20N%7D%20where%20A%3C:Union%7BReinterpretArray%7BT,N,S,A%7D%20where%20S%20where%20A%3C:Union%7BSubArray%7BT,N,A,I,true%7D%20where%20I%3C:Union%7BTuple%7BVararg%7BReal,N%7D%20where%20N%7D,%20Tuple%7BAbstractUnitRange,Vararg%7BAny,N%7D%20where%20N%7D%7D%20where%20A%3C:DenseArray%20where%20N%20where%20T,%20DenseArray%7D%20where%20N%20where%20T,%20SubArray%7BT,N,A,I,true%7D%20where%20I%3C:Union%7BTuple%7BVararg%7BReal,N%7D%20where%20N%7D,%20Tuple%7BAbstractUnitRange,Vararg%7BAny,N%7D%20where%20N%7D%7D%20where%20A%3C:DenseArray%20where%20N%20where%20T,%20DenseArray%7D,%20SubArray%7B#s36,2,A,I,L%7D%20where%20L%20where%20I%3C:Tuple%7BVararg%7BUnion%7BInt64,%20AbstractRange%7BInt64%7D,%20AbstractCartesianIndex%7D,N%7D%20where%20N%7D%20where%20A%3C:Union%7BReinterpretArray%7BT,N,S,A%7D%20where%20S%20where%20A%3C:Union%7BSubArray%7BT,N,A,I,true%7D%20where%20I%3C:Union%7BTuple%7BVararg%7BReal,N%7D%20where%20N%7D,%20Tuple%7BAbstractUnitRange,Vararg%7BAny,N%7D%20where%20N%7D%7D%20where%20A%3C:DenseArray%20where%20N%20where%20T,%20DenseArray%7D%20where%20N%20where%20T,%20ReshapedArray%7BT,N,A,MI%7D%20where%20MI%3C:Tuple%7BVararg%7BSignedMultiplicativeInverse%7BInt64%7D,N%7D%20where%20N%7D%20where%20A%3C:Union%7BReinterpretArray%7BT,N,S,A%7D%20where%20S%20where%20A%3C:Union%7BSubArray%7BT,N,A,I,true%7D%20where%20I%3C:Union%7BTuple%7BVararg%7BReal,N%7D%20where%20N%7D,%20Tuple%7BAbstractUnitRange,Vararg%7BAny,N%7D%20where%20N%7D%7D%20where%20A%3C:DenseArray%20where%20N%20where%20T,%20DenseArray%7D%20where%20N%20where%20T,%20SubArray%7BT,N,A,I,true%7D%20where%20I%3C:Union%7BTuple%7BVararg%7BReal,N%7D%20where%20N%7D,%20Tuple%7BAbstractUnitRange,Vararg%7BAny,N%7D%20where%20N%7D%7D%20where%20A%3C:DenseArray%20where%20N%20where%20T,%20DenseArray%7D%20where%20N%20where%20T,%20DenseArray%7D%7D%20where%20#s36%3C:Union%7BComplex%7BFloat32%7D,%20Complex%7BFloat64%7D,%20Float32,%20Float64%7D%7D))? This type of linear algebra occurs in probabilistic programming of a likelihood based on a Kalman filter type of computation.

The issue with algebra functions is that they rely on FORTRAN (BLAS, LAPACK). There are some options for pure Julia linear algebra but they are quite limited in scope. It seems that the way forward for these functions is a custom derivative. For one, this eliminates the issue of autodiff through FORTRAN code. Furthermore, the result could run faster, and potentially also be more accurate.

One solution is to define a custom Jacobian. This would be great and would help me to move forward. However, potentially there is a better solution based on the [Fréchet derivative](https://en.wikipedia.org/wiki/Fr%C3%A9chet_derivative).

Some examples can illustrate these points using the matrix inverse `inv`. For a 40x40 matrix, I can compute the Jacobian in 80 milliseseconds using:

```julia
using LinearAlgebra
using ForwardDiff

n = 40
A = randn(rng,n,n)
ForwardDiff.jacobian(inv,A)
# 79.573 ms (2015 allocations: 84.73 MiB)

```

Using the custom Jacobian:

```julia
function jacobian_inv(A)
  Ainv = inv(A)
  -kron(Ainv',Ainv)
end

```

The same Jacobian is computed in 7 milliseconds, a 10x speed-up:

```julia
jacobian_inv(A)
# 7.074 ms (10 allocations: 39.10 MiB)

```

However, the approach based on the Fréchet derivative can be more computationally efficient, as well as easier to formulate. In this approach, the derivative is not a number or a matrix, but a linear function. The chain rule is not a chain of multiplications but a chain of linear functions. The benefit is that in practice it may happen that we only want to evaluate this linear function at one given value, reducing computational load. In this sense, derivatives for vector-valued functions are quite different from scalar derivatives. As an illustration, consider the following function `ftest`:

```julia
n = 40
A = randn(rng,n,n)
B = randn(rng,n,n)
C = randn(rng,n,n)
ftest(x) = (C/(A+x[1]*B))[end,end] # = (C*inv(A+x[1]*B))[end,end]

```

The Fréchet derivative for the matrix inverse is given by:

```julia
function frechet_inv(A,Δ)
  A = lu(A)
  -(A\Δ/A)
end

```

We can now perform differentiation of `ftest` using:

```julia
xtest = 1
(C*frechet_inv(A+xtest*B,B))[end,end]
# 68.750 μs (11 allocations: 101.52 KiB)

```

which runs in 69 μs, 100x faster than the custom Jacobian! (ForwardDiff takes 141 μs for this function).

To summarize: Is there an existing package with which I can formulate a custom Jacobian, or, better yet, a custom Fréchet derivative? Besides supporting autodiff for linear algebra, this would also enable autodiff more generally for black box vector-valued functions. In absence of this, is there a developing package where I can bring up this request? Perhaps [capstan](https://github.com/JuliaDiff/Capstan.jl) is an option, although this is very early in its development cycle.

---

<div class="post-metadata">

**Author:** ![improbable22](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/improbable22/32/5464_2.png) [@improbable22](https://discourse.julialang.org/u/improbable22)\
**Post date:** [April 6, 2019, 3:57am UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/2 "2019-04-06T03:57:31Z")

</div>

What you are calling a Fréchet derivative appears to be something very close to how backward-mode AD works? Most packages let you define custom gradients, with `@grad` in Flux you provide exactly the backward function (which gets sewn into a chain).

---

<div class="post-metadata">

**Author:** ![sdewaele](https://avatars.discourse-cdn.com/v4/letter/s/7ab992/32.png) [@sdewaele](https://discourse.julialang.org/u/sdewaele)\
**Post date:** [April 6, 2019, 3:30pm UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/3 "2019-04-06T15:30:24Z")

</div>

Thanks for your comments! It is my understanding that Flux’s `@grad` only allows definition of [custom gradients](https://fluxml.ai/Flux.jl/stable/internals/tracker/#Custom-Gradients-1), i.e. derivatives for functions mapping vectors to _scalars_ ( ℝⁿ → ℝ). Conversely, the linear algebra operations require custom derivatives (Jacobian/Fréchet) for a mapping from vectors to _vectors_ ( ℝⁿ → ℝᵐ).

---

<div class="post-metadata">

**Author:** ![jlperla](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jlperla/32/34332_2.png) [@jlperla](https://discourse.julialang.org/u/jlperla)\
**Post date:** [April 6, 2019, 4:10pm UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/4 "2019-04-06T16:10:57Z")

</div>

My understanding is that the next generation of the Flux.Tracker, Zygote.jl, has recently become stable enough to start using. It is somewhat bleeding edge but certainly the future here (ie Capstan.jl was a proof of concept for the techniques used in Zygote.jl that may never be exist past that point).

You can look at the docs, [Custom Adjoints · Zygote](http://fluxml.ai/Zygote.jl/dev/adjoints/) but from first glance it looked like it could handle what you are trying to do. If you end up using it, I bet lots of people here would love to hear your thoughts.

---

<div class="post-metadata">

**Author:** ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)\
**Post date:** [April 6, 2019, 4:41pm UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/5 "2019-04-06T16:41:07Z")

</div>

Custom gradients only require that the last function in a chain is R^n \rightarrow R, intermediate functions can be any R^n \rightarrow R^m. Usually in backpropagation you don’t compute Jacobians explicitly, but instead use functions similar to your `frechet_inv` to push derivatives back one step at a time.

---

<div class="post-metadata">

**Author:** ![sdewaele](https://avatars.discourse-cdn.com/v4/letter/s/7ab992/32.png) [@sdewaele](https://discourse.julialang.org/u/sdewaele)\
**Post date:** [April 6, 2019, 11:09pm UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/6 "2019-04-06T23:09:02Z")

</div>

@dfdx, I think what we need for linear algebra is a (Frechet) derivative for a _vector_ - valued function, so many of the existing custom gradient options don’t seem to fit the bill. @jlperla thanks for the pointer to Zygote. This is tantalizingly close to what I am looking for, and I believe it is very suitable to include this functionality. The [pullbacks](http://fluxml.ai/Zygote.jl/dev/adjoints/#Pullbacks-1) have exactly the form of the Frechet derivatives.

I have implemented my test case using Zygote, see below. Unfortunately, the custom adjoint that I define for the matrix inverse (desired method) does not give the correct result at this point.

Any further thoughts?

```julia
using LinearAlgebra
using Random
using BenchmarkTools
using Test
using Zygote
using Zygote: @adjoint

function frechet_inv(A,Δ)
  A = lu(A)
  -(A\Δ/A)
end

ftest_inv(x) = (C*inv(A+x*B))[end,end]

# Desired method: define custom adjoint for the inverse
myinv(A) = inv(A)
@adjoint myinv(A) = myinv(A),cb -> (frechet_inv(A,cb),)
ftest_myinv(x) = (C*myinv(A+x*B))[end,end]

# Impractical method - just to show the correctness of the Frechet derivative
ftest_inv_custom(x) = (C*inv(A+x*B))[end,end]
@adjoint ftest_inv_custom(x) = ftest_inv_custom(x), cb -> ((C*frechet_inv(A+x*B,B))[end,end]*cb,)

Zygote.refresh()

rng = Random.MersenneTwister(23445)
n = 40
A = randn(rng,n,n)
B = randn(rng,n,n)
C = randn(rng,n,n)

@show size(A)
xtest = 2.3

@show ftest_inv(xtest)
# size(A) = (40, 40)

# Zygote
println("Zygote inv")
@show Zygote.gradient(ftest_inv,xtest)
y,back = Zygote.forward(ftest_inv,xtest)
@btime Zygote.forward($ftest_inv,$xtest)
# 71.220 μs (85 allocations: 73.50 KiB)

println("Zygote ftest_inv_custom - impractical")
@show Zygote.gradient(ftest_inv_custom,xtest)
y2,back_custom = Zygote.forward(ftest_inv_custom,xtest)
display(@test back_custom(1)[1]≈back(1)[1])
# Test passed

@btime Zygote.forward($ftest_inv_custom,$xtest)
# 61.345 μs (11 allocations: 71.14 KiB)

println("Zygote myinv - desired method")
@show Zygote.gradient(ftest_myinv,xtest)
y2,myback = Zygote.forward(ftest_myinv,xtest)

@test back(1)[1]≈myback(1)[1]
# Test failed
# Evaluated: -2302.6939568625844 ≈ 58.122801062141754

```

```julia

```

---

<div class="post-metadata">

**Author:** ![dfdx](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dfdx/32/120_2.png) [@dfdx](https://discourse.julialang.org/u/dfdx)\
**Post date:** [April 6, 2019, 11:33pm UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/7 "2019-04-06T23:33:13Z")

</div>

> [@sdewaele](#):
>
> @dfdx, I think what we need for linear algebra is a (Frechet) derivative for a _vector_ - valued function, so many of the existing custom gradient options don’t seem to fit the bill.

Note that your `ftest_inv()` _is_ scalar-valued. The form of other functions that `ftest_inv()` depends on doesn’t matter for custom gradients.

---

<div class="post-metadata">

**Author:** ![kolia](https://avatars.discourse-cdn.com/v4/letter/k/c68b51/32.png) [@kolia](https://discourse.julialang.org/u/kolia)\
**Post date:** [April 7, 2019, 12:31pm UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/8 "2019-04-07T12:31:06Z")

</div>

Flux and other autodiff packages most definitely let you define gradients of vector valued functions. You wouldn’t get very far without them, many of the functions you want to differentiate are vector or higher order tensor valued.

Take a look at the examples in [https://github.com/FluxML/Flux.jl/blob/master/docs/src/internals/tracker.md#custom-gradients](https://github.com/FluxML/Flux.jl/blob/master/docs/src/internals/tracker.md#custom-gradients). The minus and multiply functions whose grad are defined return general tensors/arrays. In particular, the \delta in those definitions is not restricted to being a scalar.

---

<div class="post-metadata">

**Author:** ![sdewaele](https://avatars.discourse-cdn.com/v4/letter/s/7ab992/32.png) [@sdewaele](https://discourse.julialang.org/u/sdewaele)\
**Post date:** [April 8, 2019, 3:12am UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/9 "2019-04-08T03:12:05Z")

</div>

Thanks! You are right. After your pointer I looked further into `@grad` and actually found it’s code for the custom gradient for the matrix inverse [here](https://github.com/FluxML/Tracker.jl/blob/4b64e81d56f6700964534bbaf4e5a252ff7260e8/src/lib/array.jl#L266). Fortunately, it is Julia so we can easily read the source code 🙂 . I should have done it earlier!

What confused me is (a) the name “gradient” ([definition](https://en.wikipedia.org/wiki/Gradient#Definition)) and (b) the fact that what `@grad` expects is almost, but not quite, the Frechet derivative. This also derailed my example above with Zygote. For easy reference, here is the correct custom gradient for the matrix inverse:

```julia
inv(A::TrackedArray) = Tracker.track(inv, A)
@grad function inv(A)
    return inv(Tracker.data(A)), function (Δ)
        Ainv = inv(A)
        ∇A = - Ainv' * Δ * Ainv'
        return (∇A, )
    end
end

```

---

<div class="post-metadata">

**Author:** ![antoine-levitt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/antoine-levitt/32/4008_2.png) [@antoine-levitt](https://discourse.julialang.org/u/antoine-levitt)\
**Post date:** [April 8, 2019, 9:13am UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/10 "2019-04-08T09:13:59Z")

</div>

> [@sdewaele](#):
>
> (b) the fact that what `@grad` expects is almost, but not quite, the Frechet derivative.

It is the transpose of the Frechet derivative, because of the way reverse-mode differentiation works. The Frechet derivative would be used directly in e.g. ForwardDiff (although I don’t see in the docs for ForwardDiff where you’re supposed to define custom derivatives)

---

<div class="post-metadata">

**Author:** ![sdewaele](https://avatars.discourse-cdn.com/v4/letter/s/7ab992/32.png) [@sdewaele](https://discourse.julialang.org/u/sdewaele)\
**Post date:** [April 9, 2019, 12:00am UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/11 "2019-04-09T00:00:06Z")

</div>

Thanks. To make it fully explicit, this is the relationship, where `F` is the Frechet derivative, and `G` is the function used by `@grad`:

```julia
G(Δ) = (F(Δ'))'

```

Okay, I can work with it, but it is not very intuitive to me… Also, it leads to more transposing in G, as the example of the inverse shows, as well as other functions [Tracker.jl/array.jl](https://github.com/FluxML/Tracker.jl/blob/4b64e81d56f6700964534bbaf4e5a252ff7260e8/src/lib/array.jl#L266).

---

<div class="post-metadata">

**Author:** ![antoine-levitt](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/antoine-levitt/32/4008_2.png) [@antoine-levitt](https://discourse.julialang.org/u/antoine-levitt)\
**Post date:** [April 9, 2019, 6:52am UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/12 "2019-04-09T06:52:54Z")

</div>

Again, that’s because flux is made for reverse mode autodiff (computing gradients). That’s the important primitive in that setting. Fréchet derivative is the important primitive in forward-mode autodiff, so if you want that you should look into ForwardDiff.

---

<div class="post-metadata">

**Author:** ![sdewaele](https://avatars.discourse-cdn.com/v4/letter/s/7ab992/32.png) [@sdewaele](https://discourse.julialang.org/u/sdewaele)\
**Post date:** [October 11, 2019, 1:33pm UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/13 "2019-10-11T13:33:51Z")

</div>

For those interested - I have raised [this issue with ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl/issues/413) to find out how to define custom derivatives.

---

<div class="post-metadata">

**Author:** ![c42f](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/c42f/32/52842_2.png) [@c42f](https://discourse.julialang.org/u/c42f)\
**Post date:** [October 11, 2019, 2:05pm UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/14 "2019-10-11T14:05:27Z")

</div>

You might also be interested in [ChainRules.jl](http://www.juliadiff.org/ChainRules.jl/dev/), which is all about defining custom derivatives in a very flexible and reusable way which can be integrated into various autodiff frameworks. There’s ongoing work to integrate it into Zygote and Nabla, for example.

---

<div class="post-metadata">

**Author:** ![oxinabox](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/oxinabox/32/206603_2.png) [@oxinabox](https://discourse.julialang.org/u/oxinabox)\
**Post date:** [October 14, 2019, 1:50pm UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/15 "2019-10-14T13:50:16Z")

</div>

ChainRules is also already built into  
the still very bleeding edge WIP [ForwardDiff2.jl](https://github.com/YingboMa/ForwardDiff2.jl)  
and in the medium term hope is to more or less put it everywhere  
that DiffRules currently is used.  
(which is more than just AD)

---

<div class="post-metadata">

**Author:** ![sdewaele](https://avatars.discourse-cdn.com/v4/letter/s/7ab992/32.png) [@sdewaele](https://discourse.julialang.org/u/sdewaele)\
**Post date:** [October 14, 2019, 10:23pm UTC](https://discourse.julialang.org/t/autodiff-custom-derivatives-of-vector-functions-for-linear-algebra/22826/16 "2019-10-14T22:23:51Z")

</div>

Interesting! I am happy to see the progress in this area.