# Zygote @adjoint with matrices
**URL:** https://discourse.julialang.org/t/zygote-adjoint-with-matrices/32188
**Category:** Machine Learning
**Created:** [December 12, 2019, 10:32am UTC](https://discourse.julialang.org/t/zygote-adjoint-with-matrices/32188 "2019-12-12T10:32:31Z")
**Posts on this page:** 8
**Page:** 1
### Author:  [@theogf](https://discourse.julialang.org/u/theogf)
#### Post date: [December 12, 2019, 10:32am UTC](https://discourse.julialang.org/t/zygote-adjoint-with-matrices/32188/1 "2019-12-12T10:32:31Z")
Hi!
I would like to use Zygote, because of the amazing possibility to tag the parameters you want to optimize and let Zygote do the rest.
However in my case I am creating a matrix (`K`) given a (potentially highly nested) list of parameters (`theta`) and passing them to a function (`f`) to get a scalar.
I have derived analytically the gradient `df/dtheta = g(dK/dtheta)` which is non-linear (and contains a share of optimization tricks) and Zygote works perfectly for differentiating `K`. My initial solution was to compute `dK/dp`, for each `theta` and pass it to `df/dtheta` but it is very inefficient/unpractical.
Now I want to write an `@adjoint` that would contain `g(K)` but I have no idea how to go about it since it’s not a jacobian-vector product anymore…
Here is a simple example with the derivations
How can I write an appropriate adjoint for this?
---
### Author:  [@simeonschaub](https://discourse.julialang.org/u/simeonschaub)
#### Post date: [December 12, 2019, 1:25pm UTC](https://discourse.julialang.org/t/zygote-adjoint-with-matrices/32188/2 "2019-12-12T13:25:40Z")
In Zygote, the pullback, which maps the previous jacobian to the new jacobian, is just an arbitrary function, so it doesn’t necessarily have to be a jacobian-vector product. It should be as easy as:
```julia
@adjoint f(K) = f(K), J -> (g(J),)
```
---
### Author:  [@theogf](https://discourse.julialang.org/u/theogf)
#### Post date: [December 12, 2019, 2:27pm UTC](https://discourse.julialang.org/t/zygote-adjoint-with-matrices/32188/3 "2019-12-12T14:27:12Z")
The problem is that when I do this `J` is a scalar.
---
### Author:  [@simeonschaub](https://discourse.julialang.org/u/simeonschaub)
#### Post date: [December 12, 2019, 2:39pm UTC](https://discourse.julialang.org/t/zygote-adjoint-with-matrices/32188/4 "2019-12-12T14:39:49Z")
You’re right, Zygote does reverse-mode differentiation, so the argument to the pullback of `f` is actually `df/df`, which is just one. In your case, I would suggest looking into forward-mode AD using `ForwardDiff` instead because it should be much more efficient for differentiating `K` and it will be easier to implement this custom adjoint for `f`.
---
### Author:  [@theogf](https://discourse.julialang.org/u/theogf)
#### Post date: [December 12, 2019, 2:57pm UTC](https://discourse.julialang.org/t/zygote-adjoint-with-matrices/32188/5 "2019-12-12T14:57:54Z")
The only problem is that I need the implicit differentiation of Zygote 😅
I need to rely on `Zygote.params`
---
### Author:  [@simeonschaub](https://discourse.julialang.org/u/simeonschaub)
#### Post date: [December 12, 2019, 3:02pm UTC](https://discourse.julialang.org/t/zygote-adjoint-with-matrices/32188/6 "2019-12-12T15:02:30Z")
You can use ForwardDiff within Zygote with the function `forwarddiff`. See also [here](https://fluxml.ai/Zygote.jl/latest/utils/#Zygote.forwarddiff) in the Zygote docs.
---
### Author:  [@theogf](https://discourse.julialang.org/u/theogf)
#### Post date: [December 12, 2019, 3:17pm UTC](https://discourse.julialang.org/t/zygote-adjoint-with-matrices/32188/7 "2019-12-12T15:17:20Z")
Thanks but this is not compatible with the `Params` approach 🙂
---
### Author:  [@theogf](https://discourse.julialang.org/u/theogf)
#### Post date: [December 14, 2019, 1:41pm UTC](https://discourse.julialang.org/t/zygote-adjoint-with-matrices/32188/8 "2019-12-14T13:41:42Z")
I finally found a work around!
Since my concern was optimizing the gradients computations (avoiding precomputed inverses etc), I simply wrote a new function whose gradient is equivalent to the one I want!
It slightly less efficient but works pretty well for now!