Add Type Checking in Forward Mode (#808)

sunxd3 · web-flow · commit 4071f3214b05 · 2025-12-02T04:43:10.000Z
* code and test * Revert "code and test" This reverts commit a9af7cb. * code and test * formatting * try fixing test error * remove confusing example * address Bruno's suggestions * typo * separate the two debug mode tests * remove forward debug mode for temp test * fix 1.10 segfault * use generated function to avoid ever trap into the compiler error * formatting * more formatting * add couple of test to mirror reverse mode * bump version
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Mooncake"
 uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.4.180"
+version = "0.4.181"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
diff --git a/src/debug_mode.jl b/src/debug_mode.jl
@@ -1,5 +1,135 @@
-# TODO: make it non-trivial. See https://github.com/chalk-lab/Mooncake.jl/issues/672
-DebugFRule(rule) = rule
+"""
+    DebugFRule(rule)
+
+Construct a callable equivalent to `rule` but with additional type checking for forward-mode
+AD. Checks:
+- Each `Dual` argument has tangent type matching `tangent_type(typeof(primal))`
+- The returned `Dual` has a correctly-typed tangent
+- Deep structural validation (array sizes, field types, etc.)
+
+Forward-mode counterpart to [`DebugRRule`](@ref).
+
+*Note:* Debug mode significantly slows execution (10-100x) and should only be used for
+diagnosing problems, not production runs.
+```
+"""
+struct DebugFRule{Trule}
+    rule::Trule
+end
+
+# Recursively copy the wrapped rule
+_copy(x::P) where {P<:DebugFRule} = P(_copy(x.rule))
+
+"""
+    (rule::DebugFRule)(x::Vararg{Dual,N}) where {N}
+
+Apply pre- and post-condition type checking. See [`DebugFRule`](@ref).
+"""
+@static if VERSION < v"1.11-"
+    # On Julia 1.10, use @generated to check types at compile time, preventing the
+    # compiler from ever seeing rule.rule(x...) with mismatched types, which would
+    # cause a segfault (JuliaLang/julia#51016).
+    @generated function (rule::DebugFRule{Trule})(x::Vararg{Dual,N}) where {Trule,N}
+        # First, check tangent type consistency for all Dual inputs at compile time.
+        # This prevents the compiler from generating code for rule.rule(x...) with
+        # mismatched Dual types (e.g., Dual{Float64,Float32} instead of Dual{Float64,Float64}).
+        for dt in x
+            P = dt.parameters[1]  # primal type
+            T = dt.parameters[2]  # tangent type
+            T_expected = tangent_type(P)
+            if T !== T_expected
+                msg = "Error in inputs to rule with input types $(Tuple{x...})"
+                return :(error($msg))
+            end
+        end
+
+        # Check primal types match rule signature
+        if Trule <: DerivedFRule && isconcretetype(Trule)
+            sig = Trule.parameters[1]      # primal_sig
+            isva = Trule.parameters[3]
+            nargs_val = Trule.parameters[4]
+
+            # Extract primal types
+            primal_types = [dt.parameters[1] for dt in x]
+
+            # Handle varargs unflattening
+            if isva
+                regular_types = primal_types[1:(nargs_val - 1)]
+                vararg_types = primal_types[nargs_val:end]
+                grouped_type = Tuple{vararg_types...}
+                final_types = [regular_types..., grouped_type]
+            else
+                final_types = primal_types
+            end
+
+            Tx = Tuple{final_types...}
+            if !(Tx <: sig)
+                msg = "Error in inputs to rule with input types $(Tuple{x...})"
+                return :(error($msg))
+            end
+        end
+
+        return quote
+            verify_dual_inputs(x)
+            y = rule.rule(x...)
+            verify_dual_output(x, y)
+            return y
+        end
+    end
+else
+    @noinline function (rule::DebugFRule)(x::Vararg{Dual,N}) where {N}
+        try
+            verify_args(rule.rule, x)
+        catch
+            error("Error in inputs to rule with input types $(_typeof(x))")
+        end
+        verify_dual_inputs(x)
+        y = rule.rule(x...)
+        verify_dual_output(x, y)
+        return y::Dual
+    end
+end
+
+@noinline function verify_dual_inputs(@nospecialize(x::Tuple))
+    try
+        for _x in x
+            _x isa Dual || error("Expected Dual, got $(typeof(_x))")
+            verify_dual_value(_x)
+        end
+    catch e
+        error("Error in inputs to rule with input types $(_typeof(x))")
+    end
+end
+
+@noinline function verify_dual_output(@nospecialize(x), @nospecialize(y))
+    try
+        y isa Dual || error("frule!! must return a Dual, got $(typeof(y))")
+        verify_dual_value(y)
+    catch e
+        error("Error in outputs of rule with input types $(_typeof(x))")
+    end
+end
+
+@noinline function verify_dual_value(d::Dual{P,T}) where {P,T}
+    # Fast path: type-level check using the Dual type parameters to enforce T == tangent_type(P)
+    T_expected = tangent_type(P)
+    if T !== T_expected
+        throw(
+            InvalidFDataException(
+                "Dual tangent type mismatch: primal $P requires tangent type " *
+                "$T_expected, but got $T",
+            ),
+        )
+    end
+
+    # Slow path: deep structural validation
+    p, t = primal(d), tangent(d)
+    # We validate fdata and rdata separately so these helpers stay in sync with reverse-mode checks.
+    verify_fdata_value(p, fdata(t))
+    verify_rdata_value(p, rdata(t))
+
+    return nothing
+end
 
 """
     DebugPullback(pb, y, x)
diff --git a/src/interpreter/forward_mode.jl b/src/interpreter/forward_mode.jl
@@ -89,6 +89,18 @@ function _copy(x::P) where {P<:DerivedFRule}
     return P(replace_captures(x.fwd_oc, _copy(x.fwd_oc.oc.captures)))
 end
 
+_isva(::DerivedFRule{P,T,isva,nargs}) where {P,T,isva,nargs} = isva
+_nargs(::DerivedFRule{P,T,isva,nargs}) where {P,T,isva,nargs} = nargs
+
+# Extends functionality defined in debug_mode.jl.
+function verify_args(r::DerivedFRule{sig}, x) where {sig}
+    Tx = Tuple{
+        map(_typeof ∘ primal, __unflatten_dual_varargs(_isva(r), x, Val(_nargs(r))))...
+    }
+    Tx <: sig && return nothing
+    throw(ArgumentError("Arguments with sig $Tx do not subtype rule signature, $sig"))
+end
+
 """
     __unflatten_dual_varargs(isva::Bool, args, ::Val{nargs}) where {nargs}
 
diff --git a/src/test_utils.jl b/src/test_utils.jl
@@ -96,6 +96,7 @@ using Mooncake:
     MutableTangent,
     frule!!,
     rrule!!,
+    DebugFRule,
     build_rrule,
     tangent_type,
     zero_tangent,
@@ -944,8 +945,8 @@ function test_rule(
     frule = test_fwd ? build_frule(fwd_interp, sig; debug_mode) : missing
     rrule = test_rvs ? build_rrule(rvs_interp, sig; debug_mode) : missing
 
-    # If something is primitive, then the rule should be `rrule!!`.
-    test_fwd && is_primitive && @test frule == frule!!
+    # If something is primitive, then the rule should be `frule!!` or `rrule!!`.
+    test_fwd && is_primitive && @test frule == (debug_mode ? DebugFRule(frule!!) : frule!!)
     test_rvs && is_primitive && @test rrule == (debug_mode ? DebugRRule(rrule!!) : rrule!!)
 
     # Generate random tangents for anything that is not already a CoDual.
diff --git a/test/debug_mode.jl b/test/debug_mode.jl
@@ -1,44 +1,142 @@
 @testset "debug_mode" begin
+    @testset "reverse debug mode" begin
+        # Unless we explicitly check that the arguments are of the type as expected by the rule,
+        # this will segfault.
+        @testset "argument checking" begin
+            f = x -> 5x
+            rule = build_rrule(f, 5.0; debug_mode=true)
+            @test_throws ErrorException rule(zero_fcodual(f), CoDual(0.0f0, 1.0f0))
+        end
 
-    # Unless we explicitly check that the arguments are of the type as expected by the rule,
-    # this will segfault.
-    @testset "argument checking" begin
-        f = x -> 5x
-        rule = build_rrule(f, 5.0; debug_mode=true)
-        @test_throws ErrorException rule(zero_fcodual(f), CoDual(0.0f0, 1.0f0))
+        # Forwards-pass tests.
+        x = (CoDual(sin, NoTangent()), CoDual(5.0, NoFData()))
+        @test_throws(ErrorException, Mooncake.DebugRRule(rrule!!)(x...))
+        x = (CoDual(sin, NoFData()), CoDual(5.0, NoFData()))
+        @test_throws(
+            ErrorException,
+            Mooncake.DebugRRule((x...,) -> (CoDual(1.0, 0.0), nothing))(x...)
+        )
+
+        # Basic type checking.
+        x = (CoDual(size, NoFData()), CoDual(randn(10), randn(Float16, 11)))
+        @test_throws ErrorException Mooncake.DebugRRule(rrule!!)(x...)
+
+        # Element type checking. Abstractly typed-elements prevent determining incorrectness
+        # just by looking at the array.
+        x = (
+            CoDual(size, NoFData()),
+            CoDual(Any[rand() for _ in 1:10], Any[rand(Float16) for _ in 1:10]),
+        )
+        @test_throws ErrorException Mooncake.DebugRRule(rrule!!)(x...)
+
+        # Test that bad rdata is caught as a pre-condition.
+        y, pb!! = Mooncake.DebugRRule(rrule!!)(zero_fcodual(sin), zero_fcodual(5.0))
+        @test_throws(InvalidRDataException, pb!!(5))
+
+        # Test that bad rdata is caught as a post-condition.
+        rule_with_bad_pb(x::CoDual{Float64}) = x, dy -> (5,) # returns the wrong type
+        y, pb!! = Mooncake.DebugRRule(rule_with_bad_pb)(zero_fcodual(5.0))
+        @test_throws InvalidRDataException pb!!(1.0)
+
+        # Test that bad rdata is caught as a post-condition.
+        rule_with_bad_pb_length(x::CoDual{Float64}) = x, dy -> (5, 5.0) # returns the wrong type
+        y, pb!! = Mooncake.DebugRRule(rule_with_bad_pb_length)(zero_fcodual(5.0))
+        @test_throws ErrorException pb!!(1.0)
     end
 
-    # Forwards-pass tests.
-    x = (CoDual(sin, NoTangent()), CoDual(5.0, NoFData()))
-    @test_throws(ErrorException, Mooncake.DebugRRule(rrule!!)(x...))
-    x = (CoDual(sin, NoFData()), CoDual(5.0, NoFData()))
-    @test_throws(
-        ErrorException, Mooncake.DebugRRule((x...,) -> (CoDual(1.0, 0.0), nothing))(x...)
-    )
-
-    # Basic type checking.
-    x = (CoDual(size, NoFData()), CoDual(randn(10), randn(Float16, 11)))
-    @test_throws ErrorException Mooncake.DebugRRule(rrule!!)(x...)
-
-    # Element type checking. Abstractly typed-elements prevent determining incorrectness
-    # just by looking at the array.
-    x = (
-        CoDual(size, NoFData()),
-        CoDual(Any[rand() for _ in 1:10], Any[rand(Float16) for _ in 1:10]),
-    )
-    @test_throws ErrorException Mooncake.DebugRRule(rrule!!)(x...)
-
-    # Test that bad rdata is caught as a pre-condition.
-    y, pb!! = Mooncake.DebugRRule(rrule!!)(zero_fcodual(sin), zero_fcodual(5.0))
-    @test_throws(InvalidRDataException, pb!!(5))
-
-    # Test that bad rdata is caught as a post-condition.
-    rule_with_bad_pb(x::CoDual{Float64}) = x, dy -> (5,) # returns the wrong type
-    y, pb!! = Mooncake.DebugRRule(rule_with_bad_pb)(zero_fcodual(5.0))
-    @test_throws InvalidRDataException pb!!(1.0)
-
-    # Test that bad rdata is caught as a post-condition.
-    rule_with_bad_pb_length(x::CoDual{Float64}) = x, dy -> (5, 5.0) # returns the wrong type
-    y, pb!! = Mooncake.DebugRRule(rule_with_bad_pb_length)(zero_fcodual(5.0))
-    @test_throws ErrorException pb!!(1.0)
+    @testset "forward debug mode" begin
+        @testset "argument checking" begin
+            f = x -> 5x
+            rule = Mooncake.build_frule(zero_dual(f), 5.0; debug_mode=true)
+            @test_throws ErrorException rule(
+                zero_dual(f), Mooncake.Dual(Float32(5.0), Float32(1.0))
+            )
+        end
+
+        @testset "valid inputs pass" begin
+            # Single argument - use Float64, not π which has NoTangent
+            rule = Mooncake.build_frule(zero_dual(sin), 0.0; debug_mode=true)
+            @test rule(zero_dual(sin), Mooncake.Dual(3.14, 1.0)) isa Mooncake.Dual
+
+            # Multiple arguments
+            f_mul(x, y) = x * y
+            rule = Mooncake.build_frule(zero_dual(f_mul), 2.0, 3.0; debug_mode=true)
+            @test rule(
+                zero_dual(f_mul), Mooncake.Dual(2.0, 1.0), Mooncake.Dual(3.0, 0.5)
+            ) isa Mooncake.Dual
+
+            # Arrays
+            h(x) = sum(x)
+            rule = Mooncake.build_frule(zero_dual(h), randn(5); debug_mode=true)
+            @test rule(zero_dual(h), Mooncake.Dual(randn(5), randn(5))) isa Mooncake.Dual
+
+            # NoTangent (non-differentiable)
+            rule = Mooncake.build_frule(zero_dual(identity), 5; debug_mode=true)
+            @test rule(zero_dual(identity), Mooncake.Dual(5, NoTangent())) isa Mooncake.Dual
+        end
+
+        @testset "size mismatch detected" begin
+            rule = Mooncake.build_frule(zero_dual(size), randn(10); debug_mode=true)
+            @test_throws ErrorException rule(
+                zero_dual(size), Mooncake.Dual(randn(11), randn(10))
+            )
+        end
+
+        @testset "element type mismatch detected" begin
+            rule = Mooncake.build_frule(zero_dual(identity), Any[1.0]; debug_mode=true)
+            @test_throws ErrorException rule(
+                zero_dual(identity), Mooncake.Dual(Any[1.0], Any[Float16(1.0)])
+            )
+        end
+
+        @testset "scalar type mismatch detected" begin
+            rule = Mooncake.build_frule(zero_dual(identity), 1.0; debug_mode=true)
+            @test_throws ErrorException rule(
+                zero_dual(identity), Mooncake.Dual(1.0, Float32(1.0))
+            )
+        end
+
+        @testset "container type mismatch detected" begin
+            rule = Mooncake.build_frule(zero_dual(identity), (1.0, 2.0); debug_mode=true)
+            @test_throws ErrorException rule(
+                zero_dual(identity), Mooncake.Dual((1.0, 2.0), [1.0, 2.0])
+            )
+        end
+
+        @testset "output tangent type mismatch detected" begin
+            # Rule that returns wrong tangent type in output
+            bad_rule = Mooncake.DebugFRule((x...,) -> Mooncake.Dual(1.0, Float32(0.0)))
+            @test_throws ErrorException bad_rule(Mooncake.Dual(5.0, 1.0))
+        end
+
+        @testset "error messages include type info" begin
+            rule = Mooncake.build_frule(zero_dual(identity), [1.0]; debug_mode=true)
+
+            try
+                rule(zero_dual(identity), Mooncake.Dual([1.0], [Float32(1.0)]))
+                @test false  # Expected ErrorException but none was thrown
+            catch e
+                msg = sprint(showerror, e)
+                @test occursin("input types", msg)
+                @test occursin("Float", msg)  # Type info present
+            end
+        end
+
+        @testset "integration with test_rule" begin
+            # Test basic case - test_rule expects primal functions, not Duals
+            Mooncake.TestUtils.test_rule(
+                sr(123456), sin, 1.0; mode=ForwardMode, debug_mode=true, perf_flag=:none
+            )
+
+            # Test with array
+            Mooncake.TestUtils.test_rule(
+                sr(123456),
+                sum,
+                randn(5);
+                mode=ForwardMode,
+                debug_mode=true,
+                perf_flag=:none,
+            )
+        end
+    end
 end
