Stream: helpdesk (published)

Topic: ✔ Enzyme.jl Forward vs Reverse Mode


view this post on Zulip Alec (Feb 06 2024 at 02:13):

Two questions which may be related:

  1. Why does forward mode return a tuple while reverse mode return an array?
  2. How do I get the right derivative [2,-2,4] for the input array with autodiff(Forward)?
julia> using Enzyme

julia> function sum_of_squares(x)
               sum(x.^2)
       end
sum_of_squares (generic function with 1 method)

julia> gradient(Reverse,sum_of_squares,[1.,-1.,2.]),
       gradient(Forward,sum_of_squares,[1.,-1.,2.])
([2.0, -2.0, 4.0], (2.0, -2.0, 4.0))

julia> let
               d1 = zeros(3)
               d2 = zeros(3)
               a = Enzyme.autodiff(Reverse,sum_of_squares,Duplicated([1.,-1.,2.],d1))
               b = Enzyme.autodiff(Forward,sum_of_squares,Duplicated,Duplicated([1.,-1.,2.],d2))
               d1,d2, a, b
       end
([2.0, -2.0, 4.0], [0.0, 0.0, 0.0], ((nothing,),), (6.0, 0.0))

view this post on Zulip Simone Carlo Surace (Feb 22 2024 at 21:22):

When you use forward mode you inject a tangent at the input (in your case a zero array) and get out the value and velocity of the output. To get the gradient you have to repeat that for each unit vector.

view this post on Zulip Alec (Mar 06 2024 at 04:56):

Thank you, that helped enough to get the answers I was looking for.

To get the right derivative in b, d2 needs to be seeded with the initial dual number components of one:

julia> let
           d1 = zeros(3)
           d2 = [1.0, 1.0, 1.0]  # Seed vector for forward mode
           a = Enzyme.autodiff(Reverse, sum_of_squares, Duplicated([1.0, -1.0, 2.0], d1))
           b = Enzyme.autodiff(Forward, sum_of_squares, Duplicated, Duplicated([1.0, -1.0, 2.0], d2))
           d1, d2, a, b
       end
([2.0, -2.0, 4.0], [1.0, 1.0, 1.0], ((nothing,),), (6.0, 4.0))

And to get the components of the gradient, I need to get the component from each unit vector:

julia> let
                  d1 = zeros(3)
                  a = Enzyme.autodiff(Reverse, sum_of_squares, Duplicated([1.,-1.,2.], d1))

                  b1 = Enzyme.autodiff(Forward, sum_of_squares, Duplicated([1.,-1.,2.], [1.0, 0.0, 0.0]))
                  b2 = Enzyme.autodiff(Forward, sum_of_squares, Duplicated([1.,-1.,2.], [0.0, 1.0, 0.0]))
                  b3 = Enzyme.autodiff(Forward, sum_of_squares, Duplicated([1.,-1.,2.], [0.0, 0.0, 1.0]))

                  d1, a, [b1, b2, b3]
              end
([2.0, -2.0, 4.0], ((nothing,),), [(2.0,), (-2.0,), (4.0,)])

view this post on Zulip Notification Bot (Mar 06 2024 at 04:56):

Alec has marked this topic as resolved.


Last updated: Dec 28 2024 at 04:38 UTC