Tangent stripping functionality

Project.toml

@@ -1,7 +1,7 @@
 name = "Mooncake"
 uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.4.145"
+version = "0.4.146"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -83,7 +83,8 @@ JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AllocCheck", "Aqua", "BenchmarkTools", "DiffTests", "JET", "JuliaFormatter", "Pkg", "StableRNGs", "Test"]
+test = ["AllocCheck", "Aqua", "BenchmarkTools", "DiffTests", "JET", "JuliaFormatter", "Pkg", "StableRNGs", "StaticArrays", "Test"]

src/interface.jl

@@ -423,6 +423,12 @@ function __exclude_unsupported_output_internal!(y::Ptr, ::Set{UInt})
     return throw_forward_ret_type_error(y)
 end
 
+strip_tangents(t::NoTangent) = nothing
+strip_tangents(t::Union{Tangent,MutableTangent}) = strip_tangents(t.fields)
+strip_tangents(t::PossiblyUninitTangent) = strip_tangents(t.tangent)
+strip_tangents(t::Union{Tuple,NamedTuple,AbstractArray}) = map(strip_tangents, t)
+strip_tangents(t::Union{Base.IEEEFloat,Ptr}) = t
+
 """
     prepare_pullback_cache(f, x...)
 

test/interface.jl

@@ -1,4 +1,7 @@
 using Mooncake.TestUtils: count_allocs
+using Mooncake:
+    Mooncake, MutableTangent, NoTangent, PossiblyUninitTangent, Tangent, tangent_type, strip_tangents
+using StaticArrays
 
 @testset "interface" begin
     @testset "$(typeof((f, x...)))" for (ȳ, f, x...) in Any[
@@ -184,7 +187,6 @@ using Mooncake.TestUtils: count_allocs
 
         @testset "__exclude_unsupported_output , $(test_set)" for test_set in
                                                                   additional_test_set
-
             try
                 Mooncake.__exclude_unsupported_output(test_set[2])
             catch err
@@ -194,7 +196,6 @@ using Mooncake.TestUtils: count_allocs
 
         @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))
@@ -210,4 +211,81 @@ using Mooncake.TestUtils: count_allocs
             end
         end
     end
+
+    @testset "Testing strip_tangents" begin
+        struct MyStruct{A,B}
+            a::A
+            b::B
+        end
+
+        mutable struct MyStruct1{A,B}
+            a::A
+            b::B
+        end
+
+        mutable struct MyStruct2
+            x::Float64
+            MyStruct2() = new()
+            MyStruct2(x::Float64) = new(x)
+        end
+
+        test_cases = [
+            (Core.svec(1.0, 1.0), [0.0, 0.0], [0.0, 0.0]),
+            (
+                SVector(1.0),
+                (data=(0.0,),),
+                Tangent{@NamedTuple{data::Tuple{Float64}}}((data=(0.0,),)),
+            ),
+            (
+                SA[1.0, 2.0, 3.0],
+                (data=(0.0, 0.0, 0.0),),
+                Tangent{@NamedTuple{data::Tuple{Float64,Float64,Float64}}}((
+                    data=(0.0, 0.0, 0.0),
+                )),
+            ),
+            (
+                SA[1.0 2.0; 1.0 2.0],
+                (data=(0.0, 0.0, 0.0, 0.0),),
+                Tangent{@NamedTuple{data::NTuple{4,Float64}}}((data=(0.0, 0.0, 0.0, 0.0),)),
+            ),
+            (
+                MyStruct{Int64,Float64}(1, 1.0),
+                (a=nothing, b=0.0),
+                Tangent{@NamedTuple{a::NoTangent, b::Float64}}((a=NoTangent(), b=0.0)),
+            ),
+            (
+                MyStruct1{Int64,Float64}(1, 1.0),
+                (a=nothing, b=0.0),
+                MutableTangent{@NamedTuple{a::NoTangent, b::Float64}}((
+                    a=NoTangent(), b=0.0
+                )),
+            ),
+            (
+                MyStruct2(6.9518144222347e-310),
+                (x=0.0,),
+                MutableTangent{@NamedTuple{x::PossiblyUninitTangent{Float64}}}((
+                    x=PossiblyUninitTangent{Float64}(0.0),
+                )),
+            ),
+            (
+                MyStruct2(1.0),
+                (x=0.0,),
+                MutableTangent{@NamedTuple{x::PossiblyUninitTangent{Float64}}}((
+                    x=PossiblyUninitTangent{Float64}(0.0),
+                )),
+            ),
+            (1.0, 0.0, 0.0),
+            ([1.0], [0.0], [0.0]),
+            ((1.0, [1.0]), (0.0, [0.0]), (0.0, [0.0])),
+            ((1, ["hello"]), (nothing, [nothing]), (NoTangent(), [NoTangent()])),
+        ]
+
+        for (primal, resolved_tangents, tangent_data) in test_cases
+            Mooncake.__exclude_unsupported_output(primal)
+            @test strip_tangents(tangent_data) == resolved_tangents
+        end
+
+        @test strip_tangents(Ptr{Int64}(1)) == fdata(Ptr{Int64}(1))
+        @test strip_tangents(Ptr{Float64}(1)) == fdata(Ptr{Float64}(1))
+    end
 end