Syntactic sugar for vjp
#1853
Comments
|
Here's an MWE: using Enzyme
using Enzyme.EnzymeCore: ReverseModeSplit
function vjp(
rmode::ReverseModeSplit{ReturnPrimal},
dresult,
f::FA,
::Type{RA},
args::Vararg{Annotation,N},
) where {ReturnPrimal,FA<:Annotation,RA<:Annotation,N}
forward, reverse = autodiff_thunk(rmode, FA, RA, typeof.(args)...)
tape, result, shadow_result = forward(f, args...)
if RA <: Active
dinputs = only(reverse(f, args..., dresult, tape))
else
shadow_result .+= dresult # TODO: generalize beyond arrays
dinputs = only(reverse(f, args..., tape))
end
if ReturnPrimal
return (dinputs, result)
else
return (dinputs,)
end
endTestsf(x, y) = x^2 + 10y # scalar output
g(x, y) = [f(x, y)] # vector output
x, y = 3.0, 4.0
dz = -5.0
z0 = f(x, y)
dx0 = dz * 2x
dy0 = dz * 10julia> ((dx, dy), z) = vjp(ReverseSplitWithPrimal, dz, Const(f), Active, Active(x), Active(y))
((-30.0, -50.0), 49.0)
julia> z == z0 && dx == dx0 && dy == dy0
true
julia> ((dx, dy), vz) = vjp(
ReverseSplitWithPrimal, [dz], Const(g), Duplicated, Active(x), Active(y)
)
((-30.0, -50.0), [49.0])
julia> vz == [z0] && dx == dx0 && dy == dy0
true |
This was referenced Sep 23, 2024
1 task
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Uh oh!
There was an error while loading. Please reload this page.
As discussed on Slack, it is sometimes annoying to use reverse mode because you need
autodiff_thunkwhenever1.0A solution suggested by @wsmoses would be syntactic sugar for
vjp(and maybejvpbut that is less necessary causeautodiffgets the job done every time). There are a few design questions around this:dreturntogether with the return activity?Duplicatedreturns, we would have to specifydreturn = Duplicated(whatever, actual_thing_we_care_about).vjp? How do we pass several cotangents if the return value isActive?The text was updated successfully, but these errors were encountered: