Backend switching for Mooncake #768
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,46 @@ | ||
| @is_primitive MinimalCtx Tuple{DI.DifferentiateWith,<:Union{Number,AbstractArray}} | ||
|
|
||
| function Mooncake.rrule!!(dw::CoDual{<:DI.DifferentiateWith}, x::CoDual{<:Number}) | ||
| primal_func = primal(dw) | ||
|
There was a problem hiding this comment. What if there are derivatives inside There was a problem hiding this comment. im not sure i understand. dw would be a function and if derivatives are used inside i think that would be handled by the substitute backend? There was a problem hiding this comment. What I mean is that outside of this code, Mooncake may be trying to differentiate something with respect to some parameters that are stored inside |
||
| primal_x = primal(x) | ||
| (; f, backend) = primal_func | ||
| y = zero_fcodual(f(primal_x)) | ||
gdalle marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| # output is a vector, so we need to use the vector pullback | ||
AstitvaAggarwal marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| function pullback_array!!(dy::NoRData) | ||
| tx = DI.pullback(f, backend, primal_x, (fdata(y.dx),)) | ||
| return NoRData(), only(tx) | ||
|
There was a problem hiding this comment. We're not sure that There was a problem hiding this comment.
There was a problem hiding this comment. Not all callable objects are functions though. julia> struct Multiplier
a::Float64
end
julia> (m::Multiplier)(x) = m.a * x
julia> m = Multiplier(2)
Multiplier(2.0)
julia> m(3)
6.0
julia> m isa Function
falseThere was a problem hiding this comment. And I'm not even convinced the current Mooncake behavior is correct, see chalk-lab/Mooncake.jl#557 There was a problem hiding this comment. Nice catch, Mooncake treats all callable objects as a function, giving There was a problem hiding this comment. See my remark above, let's assume for simplicity that There was a problem hiding this comment. We should probably ensure that |
||
| end | ||
|
|
||
| # output is a scalar, so we can use the scalar pullback | ||
| function pullback_scalar!!(dy::Number) | ||
| tx = DI.pullback(f, backend, primal_x, (dy,)) | ||
| return NoRData(), only(tx) | ||
| end | ||
|
|
||
| return y, typeof(primal(y)) <: Number ? pullback_scalar!! : pullback_array!! | ||
| end | ||
|
|
||
| function Mooncake.rrule!!(dw::CoDual{<:DI.DifferentiateWith}, x::CoDual{<:AbstractArray}) | ||
| primal_func = primal(dw) | ||
| primal_x = primal(x) | ||
| fdata_arg = fdata(x.dx) | ||
| (; f, backend) = primal_func | ||
| y = zero_fcodual(f(primal_x)) | ||
|
|
||
| # output is a vector, so we need to use the vector pullback | ||
| function pullback_array!!(dy::NoRData) | ||
| tx = DI.pullback(f, backend, primal_x, (fdata(y.dx),)) | ||
| fdata_arg .+= only(tx) | ||
| return NoRData(), dy | ||
| end | ||
|
|
||
| # output is a scalar, so we can use the scalar pullback | ||
| function pullback_scalar!!(dy::Number) | ||
| tx = DI.pullback(f, backend, primal_x, (dy,)) | ||
| fdata_arg .+= only(tx) | ||
| return NoRData(), NoRData() | ||
| end | ||
|
|
||
| return y, typeof(primal(y)) <: Number ? pullback_scalar!! : pullback_array!! | ||
| end | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,11 +1,12 @@ | ||
| using Pkg | ||
| Pkg.add(["FiniteDiff", "ForwardDiff", "Zygote", "Mooncake"]) | ||
|
|
||
| using DifferentiationInterface, DifferentiationInterfaceTest | ||
| import DifferentiationInterfaceTest as DIT | ||
| using FiniteDiff: FiniteDiff | ||
| using ForwardDiff: ForwardDiff | ||
| using Zygote: Zygote | ||
| using Mooncake: Mooncake | ||
|
There was a problem hiding this comment. In this file we need to add tests that are specific to Mooncake. Ideally I should have done that with the other backends too. |
||
| using Test | ||
|
|
||
| LOGGING = get(ENV, "CI", "false") == "false" | ||
| @@ -24,7 +25,7 @@ function differentiatewith_scenarios() | ||
| end | ||
|
|
||
| test_differentiation( | ||
| [AutoForwardDiff(), AutoZygote(), AutoMooncake(; config=nothing)], | ||
| differentiatewith_scenarios(); | ||
| excluded=SECOND_ORDER, | ||
| logging=LOGGING, | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Update to
chalk-lab