Fix segfault when passing friendly tangents to "non-friendly" forward cache in 1.10 and formatting (#983)

Author: gdalle

src/interface.jl

@@ -690,8 +690,10 @@ Returns a `Dual` containing the result of applying forward-mode AD to compute th
 derivative of `primal(f)` at the primal values in `x` in the direction of the tangent values
 in `f` and `x`.
 """
-value_and_derivative!!(cache::ForwardCache, fx::Vararg{Dual,N}) where {N} =
-    cache.rule(fx...)  # TODO: handle friendly tangents for the output here?
+function value_and_derivative!!(cache::ForwardCache, fx::Vararg{Dual,N}) where {N}
+    # TODO: check Dual coherence here like we do below?
+    return cache.rule(fx...)
+end
 
 """
     value_and_derivative!!(cache::ForwardCache, (f, df), (x, dx), ...)
@@ -707,7 +709,7 @@ Tuples are used as inputs and outputs instead of `Dual` numbers to accommodate t
     `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}
+    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)
@@ -725,6 +727,11 @@ function value_and_derivative!!(
     end
 
     input_duals = map(Dual, input_primals, input_tangents)
+
+    if !friendly_tangents  # in friendly mode, conversion should ensure tangent coherence
+        error_if_incorrect_dual_types(input_duals...)
+    end
+
     output = cache.rule(input_duals...)
     output_primal = primal(output)
     output_tangent = tangent(output)

src/tangents/dual.jl

@@ -49,6 +49,20 @@ Check that the type of `tangent(x)` is the tangent type of the type of `primal(x
 """
 verify_dual_type(x::Dual) = tangent_type(typeof(primal(x))) == typeof(tangent(x))
 
+function error_if_incorrect_dual_types(duals::Vararg{Dual,N}) where {N}
+    correct_types = map(verify_dual_type, duals)
+    if !all(correct_types)
+        primals = map(primal, duals)
+        tangents = map(tangent, duals)
+        throw(ArgumentError("""
+        Tangent types do not match primal types:
+          - primal types:           $(map(typeof, primals))
+          - provided tangent types: $(map(typeof, tangents))
+          - required tangent types: $(map(tangent_type, map(typeof, primals)))
+        """))
+    end
+end
+
 @inline uninit_dual(x::P) where {P} = Dual(x, uninit_tangent(x))
 
 # Always sharpen the first thing if it's a type so static dispatch remains possible.

test/interface.jl

@@ -88,7 +88,7 @@ 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] 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(
@@ -101,7 +101,7 @@ end
             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] 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)
@@ -271,6 +271,7 @@ end
 
         @testset "__exclude_unsupported_output , $(test_set)" for test_set in
                                                                   additional_test_set
+
             try
                 Mooncake.__exclude_unsupported_output(test_set[2])
             catch err
@@ -280,6 +281,7 @@ 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))
@@ -351,15 +353,12 @@ end
             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
+            @test_throws ArgumentError Mooncake.value_and_derivative!!(
+                cache_sp_unfriendly, zip(fx_sp, dfx_sp)...
+            )
+            @test_throws "Tangent types do not match primal types:" Mooncake.value_and_derivative!!(
+                cache_sp_unfriendly, zip(fx_sp, dfx_sp)...
+            )
         end
     end