# Zygote AD Jacobian-vector product

**URL:** <https://discourse.julialang.org/t/zygote-ad-jacobian-vector-product/40857>\
**Category:** Specific Domains\
**Tags:** differentiation, zygote\
**Created:** [June 6, 2020, 10:39am UTC](https://discourse.julialang.org/t/zygote-ad-jacobian-vector-product/40857 "2020-06-06T10:39:00Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![pavanakumar](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pavanakumar/32/22096_2.png) [@pavanakumar](https://discourse.julialang.org/u/pavanakumar)\
**Post date:** [June 6, 2020, 10:39am UTC](https://discourse.julialang.org/t/zygote-ad-jacobian-vector-product/40857/1 "2020-06-06T10:39:00Z")

</div>

Hi

I come from the Tapenade source transformation community and use Fortran as my primary language. I recently moved to Julia and slowly adapting to the language. I have problems understanding the algorithmic differentiation in Zygote (aka differential programming).

I like to explain my problem with a simple code example. Say I want to differentiate the following function

`h(x, y) = 3x^2 + 2x + 1 + y*x - y`

In FORTRAN it translates to

```julia
SUBROUTINE CALC_H(x, y, h)
  h = 3*x*x + 2*x + 1 + y*x - y
END SUBROUTINE CALC_H

```

in the forward mode. Most AD tool including Tapenade would return the following forward gradient code,

```julia
SUBROUTINE CALC_H_D(x, xd, y, yd, h, hd)
  hd = (3*2*x+y+2)*xd + (x-1.0)*yd
  h = 3*x*x + 2*x + 1 + y*x - y
END SUBROUTINE CALC_H_D

```

So basically the function takes not just x,y but also increments xd, yd as arguments.  
This is very useful to perform Jaobian-vector products automatically by simply  
specifying xd, yd inputs. Also the output returns the gradient in hd. We have a  
similar semantics for the reverse mode.

```julia
SUBROUTINE CALC_H_B(x, xb, y, yb, h, hb)
  xb = (6*x+y+2)*hb
  yb = (x-1.0)*hb
  hb = 0.0_8
END SUBROUTINE CALC_H_B

```

So I can set the value of `hd` and get that value multiplied inside the ad routines.

All that I could do in Zygote for the moment is

```julia
julia> using Zygote
julia> h(x, y) = 3x^2 + 2x + 1 + y*x - y
h (generic function with 1 method)
julia> gradient(h, 3.0, 5.0)
(25.0, 2.0)

```

provide only two parameters `x,y`. Essentially `hd` in the adjoint mode of Zygote  
is 1.0! This is a big limitation for my work. Is there a way to achieve what I  
want in Zygote. My knowledge of Julia and Zygote is very limited.

Also how do I specify `xd, yd` in the forward mode?

Cheers.

–  
Pavanakumar Mohanamuraly

---

<div class="post-metadata">

**Author:** ![lhnguyen-vn](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lhnguyen-vn/32/15449_2.png) [@lhnguyen-vn](https://discourse.julialang.org/u/lhnguyen-vn)\
**Post date:** [June 6, 2020, 11:11am UTC](https://discourse.julialang.org/t/zygote-ad-jacobian-vector-product/40857/2 "2020-06-06T11:11:44Z")

</div>

What you’re looking for is [`Zygote.pullback`](https://fluxml.ai/Zygote.jl/latest/adjoints/#Pullbacks-1), which returns the result and a backward-pass function to calculate vector-Jacobian products.

---

<div class="post-metadata">

**Author:** ![pavanakumar](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pavanakumar/32/22096_2.png) [@pavanakumar](https://discourse.julialang.org/u/pavanakumar)\
**Post date:** [June 6, 2020, 11:43am UTC](https://discourse.julialang.org/t/zygote-ad-jacobian-vector-product/40857/3 "2020-06-06T11:43:58Z")

</div>

I am aware of pullback. But this only gives function output and the gradient with one only one of the increments. What I am referring to here is to be able to give the increment as:

Say I have

`y = sin(x)`

I need something like the one below (only for an illustration purpose)

```julia
julia> y, xd = Zygote.pullback(sin, x=0.5, xinc=0.1, yinc=0.1);

```

should return

`y, xd = xinc + cos(x) * yinc`

This is the output typically obtained from most AD tools.

pullback is able to do this

```julia
julia> h(x, y) = 3x^2 + 2x + 1 + y*x - y
julia> y, back = Zygote.pullback(h, 3.0, 5.0)
(44.0, Zygote.var"#36#37"{typeof(∂(h))}(∂(h)))
julia> back(1.0)
(25.0, 2.0)
julia> back(2.0)
(50.0, 4.0)

```

But looks like the `xinc` addition is not being done. It should be easy to add this manually I guess. But is there a similar call for the forward mode?

---

<div class="post-metadata">

**Author:** ![pavanakumar](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pavanakumar/32/22096_2.png) [@pavanakumar](https://discourse.julialang.org/u/pavanakumar)\
**Post date:** [June 6, 2020, 1:38pm UTC](https://discourse.julialang.org/t/zygote-ad-jacobian-vector-product/40857/4 "2020-06-06T13:38:16Z")

</div>

The forward mode Jac-vec is described here in the following post

> [@Efficient Vector^T \* Jacobian matrix?](https://discourse.julialang.org/t/efficient-vector-t-jacobian-matrix/21329/2):
>
> This is exactly what reverse-mode automatic differentiation is (e.g. if f is R^n to R you’re asking for a O(1) gradient). There are a few packages doing it, but it’s more tricky than forward diff.

Many thanks for pointing me in the right direction.

---

<div class="post-metadata">

**Author:** ![npr](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/npr/32/10626_2.png) [@npr](https://discourse.julialang.org/u/npr)\
**Post date:** [June 6, 2020, 3:49pm UTC](https://discourse.julialang.org/t/zygote-ad-jacobian-vector-product/40857/5 "2020-06-06T15:49:46Z")

</div>

I’m not sure I fully understand, but perhaps the `frule` and `rrule` methods in [ChainRules.jl](https://www.juliadiff.org/ChainRulesCore.jl/dev/) are what you are looking for? As of recently (and not yet in an official release) Zygote.jl uses `ChainRules.rrule` under the hood to make use of hand-coded rules for derivatives. Also Zygote is currently reverse-mode AD only (though I believe there are plans for forward-mode in Zygote).

---

<div class="post-metadata">

**Author:** ![pavanakumar](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/pavanakumar/32/22096_2.png) [@pavanakumar](https://discourse.julialang.org/u/pavanakumar)\
**Post date:** [June 7, 2020, 7:08pm UTC](https://discourse.julialang.org/t/zygote-ad-jacobian-vector-product/40857/6 "2020-06-07T19:08:04Z")

</div>

@npr Thanks! I was not aware of ChainRules.jl and its use in Zygote.  
This feels more like a traditional AD tool to me with proper push-forward  
and pull-backward operations. I will try this out.