Zygote AD Jacobian-vector product

Specific Domains
,
1

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

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,

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.

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

2

What you’re looking for is Zygote.pullback, which returns the result and a backward-pass function to calculate vector-Jacobian products.

1 Like
3

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> 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> 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?

4

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

Many thanks for pointing me in the right direction.

1 Like
5

I’m not sure I fully understand, but perhaps the frule and rrule methods in ChainRules.jl 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).

6

@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.