Add caching and friendly tangents in forward mode (#980)

Author: gdalle

Project.toml

@@ -1,7 +1,7 @@
 name = "Mooncake"
 uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.5.0"
+version = "0.5.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -92,4 +92,15 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AllocCheck", "Aqua", "BenchmarkTools", "DiffTests", "JET", "Logging", "Pkg", "Revise", "StableRNGs", "Test"]
+test = [
+    "AllocCheck",
+    "Aqua",
+    "BenchmarkTools",
+    "DiffTests",
+    "JET",
+    "Logging",
+    "Pkg",
+    "Revise",
+    "StableRNGs",
+    "Test",
+]

src/interface.jl

@@ -658,20 +658,82 @@ function value_and_gradient!!(
     end
 end
 
+struct ForwardCache{R,IT<:Union{Nothing,Tuple},OP}
+    rule::R
+    input_tangents::IT
+    output_primal::OP
+end
+
 """
     prepare_derivative_cache(fx...; config=Mooncake.Config())
 
 Returns a cache used with [`value_and_derivative!!`](@ref). See that function for more info.
 """
-@unstable function prepare_derivative_cache(fx...; config=Config())
-    build_frule(fx...; config.debug_mode, config.silence_debug_messages)
+@unstable function prepare_derivative_cache(f, x::Vararg{Any,N}; config=Config()) where {N}
+    fx = (f, x...)
+    rule = build_frule(fx...; config.debug_mode, config.silence_debug_messages)
+
+    if config.friendly_tangents
+        y = f(x...)
+        input_tangents = map(zero_tangent, fx)
+        output_primal = _copy_output(y)
+        return ForwardCache(rule, input_tangents, output_primal)
+    else
+        return ForwardCache(rule, nothing, nothing)
+    end
 end
 
 """
-    value_and_derivative!!(rule::R, f::Dual, x::Vararg{Dual,N})
+    value_and_derivative!!(cache::ForwardCache, f::Dual, x::Vararg{Dual,N})
 
 Returns a `Dual` containing the result of applying forward-mode AD to compute the (Frechet)
 derivative of `primal(f)` at the primal values in `x` in the direction of the tangent values
 in `f` and `x`.
 """
-value_and_derivative!!(rule::R, fx::Vararg{Dual,N}) where {R,N} = rule(fx...)
+value_and_derivative!!(cache::ForwardCache, fx::Vararg{Dual,N}) where {N} =
+    cache.rule(fx...)  # TODO: handle friendly tangents for the output here?
+
+"""
+    value_and_derivative!!(cache::ForwardCache, (f, df), (x, dx), ...)
+
+Returns a tuple `(y, dy)` containing the result of applying forward-mode AD to compute the (Frechet) derivative of `primal(f)` at the primal values in `x` in the direction of the tangent values contained in `df` and `dx`.
+
+Tuples are used as inputs and outputs instead of `Dual` numbers to accommodate the case where internal Mooncake tangent types do not coincide with tangents provided by the user (in which case we translate between "friendly tangents" and internal tangents using cache storage).
+
+!!! info
+    `cache` must be the output of [`prepare_derivative_cache`](@ref), and (fields of) `f` and `x` must be of the same size and shape as those used to construct the `cache`. This is to ensure that the gradient can be written to the memory allocated when the `cache` was built.
+
+!!! warning
+    `cache` owns any mutable state returned by this function, meaning that mutable components of values returned by it will be mutated if you run this function again with different arguments. Therefore, if you need to keep the values returned by this function around over multiple calls to this function with the same `cache`, you should take a copy (using `copy` or `deepcopy`) of them before calling again.
+"""
+function value_and_derivative!!(
+    cache::ForwardCache, f::NTuple{2,Any}, x::Vararg{<:NTuple{2,Any},N}
+) where {N}
+    fx = (f, x...)  # to avoid method ambiguity
+    friendly_tangents = !isnothing(cache.input_tangents)
+
+    input_primals = map(first, fx)
+    input_friendly_tangents = map(last, fx)
+
+    # translate from friendly to native
+    if friendly_tangents
+        input_tangents = map(
+            primal_to_tangent!!, cache.input_tangents, input_friendly_tangents
+        )
+    else
+        input_tangents = input_friendly_tangents
+    end
+
+    input_duals = map(Dual, input_primals, input_tangents)
+    output = cache.rule(input_duals...)
+    output_primal = primal(output)
+    output_tangent = tangent(output)
+
+    # translate from native back to friendly
+    if friendly_tangents
+        output_friendly_tangent = tangent_to_primal!!(cache.output_primal, output_tangent)
+        return output_primal, output_friendly_tangent
+    else
+        return output_primal, output_tangent
+    end
+end

test/interface.jl

@@ -88,25 +88,25 @@ end
 
             cache = Mooncake.prepare_gradient_cache(f, x)
             v, dx = Mooncake.value_and_gradient!!(cache, f, x)
-            @test dx[2] isa Mooncake.Tangent{@NamedTuple{x1::Float64,x2::Float64}}
-            @test dx[2].fields == (; x1=2*x.x1, x2=cos(x.x2))
+            @test dx[2] isa Mooncake.Tangent{@NamedTuple{x1::Float64, x2::Float64}}
+            @test dx[2].fields == (; x1=2 * x.x1, x2=cos(x.x2))
 
             cache = Mooncake.prepare_gradient_cache(
-                f, x; config=Mooncake.Config(friendly_tangents=true)
+                f, x; config=Mooncake.Config(; friendly_tangents=true)
             )
             v, dx = Mooncake.value_and_gradient!!(cache, f, x)
             @test dx[2] isa SimplePair
-            @test dx[2] == SimplePair(2*x.x1, cos(x.x2))
+            @test dx[2] == SimplePair(2 * x.x1, cos(x.x2))
 
             rule = build_rrule(f, x)
 
             v, dx = Mooncake.value_and_gradient!!(rule, f, x)
-            @test dx[2] isa Mooncake.Tangent{@NamedTuple{x1::Float64,x2::Float64}}
-            @test dx[2].fields == (; x1=2*x.x1, x2=cos(x.x2))
+            @test dx[2] isa Mooncake.Tangent{@NamedTuple{x1::Float64, x2::Float64}}
+            @test dx[2].fields == (; x1=2 * x.x1, x2=cos(x.x2))
 
             v, dx = Mooncake.value_and_gradient!!(rule, f, x; friendly_tangents=true)
             @test dx[2] isa SimplePair
-            @test dx[2] == SimplePair(2*x.x1, cos(x.x2))
+            @test dx[2] == SimplePair(2 * x.x1, cos(x.x2))
         end
     end
     @testset "value_and_pullback!!" begin
@@ -162,7 +162,7 @@ end
             )
 
             cache = Mooncake.prepare_pullback_cache(
-                testf, x; config=Mooncake.Config(friendly_tangents=true)
+                testf, x; config=Mooncake.Config(; friendly_tangents=true)
             )
             v, pb = Mooncake.value_and_pullback!!(cache, x̄, testf, x)
             @test has_equal_data(v, SimplePair(x.x1^2 + sin(x.x2), x.x1 * x.x2))
@@ -271,7 +271,6 @@ end
 
         @testset "__exclude_unsupported_output , $(test_set)" for test_set in
                                                                   additional_test_set
-
             try
                 Mooncake.__exclude_unsupported_output(test_set[2])
             catch err
@@ -281,7 +280,6 @@ end
 
         @testset "_copy_output & _copy_to_output!!, $(test_set)" for test_set in
                                                                      additional_test_set
-
             original = test_set[2]
             try
                 if isnothing(Mooncake.__exclude_unsupported_output(original))
@@ -304,11 +302,65 @@ end
         (; debug_mode=true, silence_debug_messages=true),
     ]
         f = (x, y) -> x * y + cos(x)
-        fx = (f, 5.0, 4.0)
-        rule = Mooncake.prepare_derivative_cache(fx...; config=Mooncake.Config(; kwargs...))
-        z = Mooncake.value_and_derivative!!(rule, map(zero_dual, fx)...)
-        @test z isa Mooncake.Dual
-        @test primal(z) == f(5.0, 4.0)
+        g = (sp::SimplePair) -> SimplePair(f(sp.x1, sp.x2), 2.0)
+
+        x, y = 5.0, 4.0
+        dx, dy = 3.0, 2.0
+        fx = (f, x, y)
+        dfx = (Mooncake.zero_tangent(f), dx, dy)
+        z = f(x, y)
+        dz = dx * y + x * dy + dx * (-sin(x))
+
+        fx_sp = (g, SimplePair(x, y))
+        dfx_sp = (Mooncake.zero_tangent(g), SimplePair(dx, dy))
+        z_sp = g(SimplePair(x, y))
+
+        @testset "Simple types" begin
+            cache = Mooncake.prepare_derivative_cache(
+                fx...; config=Mooncake.Config(; kwargs...)
+            )
+
+            # legacy Dual interface
+            z_and_dz_dual = Mooncake.value_and_derivative!!(
+                cache, map(Mooncake.Dual, fx, dfx)...
+            )
+            @test z_and_dz_dual isa Mooncake.Dual
+            @test Mooncake.primal(z_and_dz_dual) == z
+            @test Mooncake.tangent(z_and_dz_dual) == dz
+
+            # new tuple interface
+            z_and_dz_tup = Mooncake.value_and_derivative!!(cache, zip(fx, dfx)...)
+            @test z_and_dz_tup isa Tuple{Float64,Float64}
+            @test first(z_and_dz_tup) == z
+            @test last(z_and_dz_tup) == dz
+        end
+
+        @testset "Structured types" begin
+            cache_sp_friendly = Mooncake.prepare_derivative_cache(
+                fx_sp...; config=Mooncake.Config(; friendly_tangents=true, kwargs...)
+            )
+            # friendly input doesn't error
+            z_and_dz_sp = Mooncake.value_and_derivative!!(
+                cache_sp_friendly, zip(fx_sp, dfx_sp)...
+            )
+            # output is friendly
+            @test z_and_dz_sp isa Tuple{SimplePair,SimplePair}
+            @test first(z_and_dz_sp) == SimplePair(z, 2.0)
+            @test last(z_and_dz_sp) == SimplePair(dz, 0.0)
+
+            cache_sp_unfriendly = Mooncake.prepare_derivative_cache(
+                fx_sp...; config=Mooncake.Config(; friendly_tangents=false, kwargs...)
+            )
+            if get(kwargs, :debug_mode, false)
+                @test_throws ErrorException Mooncake.value_and_derivative!!(
+                    cache_sp_unfriendly, zip(fx_sp, dfx_sp)...
+                )
+            else
+                @test_throws TypeError Mooncake.value_and_derivative!!(
+                    cache_sp_unfriendly, zip(fx_sp, dfx_sp)...
+                )
+            end
+        end
     end
 
     @testset "selective zeroing of cotangents" begin