Add Type Checking in Forward Mode

[sunxd3]

start to address #672

mostly conceptually copy-paste

[github-actions]

Mooncake.jl documentation for PR #808 is available at:
https://chalk-lab.github.io/Mooncake.jl/previews/PR808/

[github-actions]

Performance Ratio:
Ratio of time to compute gradient and time to compute function.
Warning: results are very approximate! See here for more context.

┌────────────────────────────┬──────────┬──────────┬─────────────┬─────────┬─────────────┬────────┐
│                      Label │   Primal │ Mooncake │ MooncakeFwd │  Zygote │ ReverseDiff │ Enzyme │
│                     String │   String │   String │      String │  String │      String │ String │
├────────────────────────────┼──────────┼──────────┼─────────────┼─────────┼─────────────┼────────┤
│                   sum_1000 │ 100.0 ns │      1.8 │         1.9 │     1.1 │        5.61 │   8.21 │
│                  _sum_1000 │ 941.0 ns │     6.62 │        1.01 │  1460.0 │        34.4 │   1.09 │
│               sum_sin_1000 │  6.55 μs │     2.54 │        1.39 │    1.68 │        10.8 │   2.21 │
│              _sum_sin_1000 │  5.23 μs │     2.96 │        2.18 │   259.0 │        13.1 │   2.47 │
│                   kron_sum │ 259.0 μs │     50.2 │        3.26 │    6.15 │       254.0 │   15.0 │
│              kron_view_sum │ 315.0 μs │     43.3 │        3.56 │    11.3 │       245.0 │   6.99 │
│      naive_map_sin_cos_exp │  2.15 μs │     2.32 │         1.4 │ missing │        7.09 │   2.35 │
│            map_sin_cos_exp │  2.11 μs │      2.7 │        1.46 │    1.58 │        6.13 │   2.94 │
│      broadcast_sin_cos_exp │  2.28 μs │     2.38 │        1.39 │    2.28 │        1.46 │   2.25 │
│                 simple_mlp │ 200.0 μs │      6.2 │         2.9 │     1.8 │        10.8 │   3.37 │
│                     gp_lml │ 244.0 μs │      8.5 │        2.14 │    3.72 │     missing │   5.04 │
│ turing_broadcast_benchmark │  1.76 ms │     5.01 │        3.46 │ missing │        28.3 │   2.72 │
│         large_single_block │ 380.0 ns │     4.53 │        2.03 │  4560.0 │        32.5 │   2.24 │
└────────────────────────────┴──────────┴──────────┴─────────────┴─────────┴─────────────┴────────┘

[codecov]

Codecov Report

❌ Patch coverage is 3.57143% with 27 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/debug_mode.jl	0.00%	27 Missing ⚠️

📢 Thoughts on this report? Let us know!

[yebai]

Thanks, @sunxd3 — I’ll take a look at this soon. In the meantime, could you check whether the debugging mode helps us identify the cause of #777 / #632 (comment)

[Copilot]

Pull Request Overview

This PR adds type checking support for forward-mode automatic differentiation (AD) in debug mode, addressing issue #672. It implements DebugFRule as a forward-mode counterpart to the existing DebugRRule, providing runtime validation of tangent types in forward-mode AD operations.

Key changes:

• Implementation of DebugFRule struct and validation logic for forward-mode AD
• Comprehensive test suite covering various input types and error conditions
• Integration with existing test infrastructure

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
src/debug_mode.jl	Implements DebugFRule with type checking for Dual inputs/outputs and structural validation
src/test_utils.jl	Updates test_rule to use DebugFRule wrapper when debug_mode=true for forward-mode primitives
test/debug_mode.jl	Adds comprehensive test suite for forward-mode debug functionality covering valid inputs, type mismatches, and integration tests

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[sunxd3]

thanks Bruno

[Technici4n]

I wonder why this segfaults on 1.10 😓

[sunxd3]

Mooncake.jl/test/debug_mode.jl

Lines 84 to 89 in e4b6576

	@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

triggered the segfault.

I think they related to JuliaLang/julia#51016. Essentially, when DebugFRule receives a Dual{Float64, Float32} but the inner rule was compiled for Dual{Float64, Float64}, the compiler segfaults while trying to specialize the call.

Not certain if Base.inferencebarrier is the best solution here, but it seems to prevent the crash.

[Technici4n]

Is this maybe caused by the lack of verify_args?

[sunxd3]

I think adding verify_args alone wouldn't solve it -- it won't prevent the Julia compiler from trying to do the type inference.

I added verify_args similar to the reverse_mode one, but also had to wrap the whole debug rule function in generated block so that we can properly intercept the error.

[Technici4n]

Since this is not needed for reverse debug mode, I am a bit skeptical that it's required? 😄

[sunxd3]

it's tricky, reverse debug mode was not testing the case where the segfault appear for the forward debug mode

Mooncake.jl/test/debug_mode.jl

Lines 3 to 9 in bcfa006

	# 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

but e.g., same segfault will happen with

julia +lts --project -e '
      using Mooncake
      using Mooncake: CoDual, zero_fcodual, build_rrule

      f = x -> 5x
      rule = build_rrule(f, 5.0; debug_mode=true)

      rule(zero_fcodual(f), CoDual(5.0, 1.0))
      '

[Technici4n]

Ah I see. Funny how simple PRs turn into a rabbit hole sometimes... 😅

[sunxd3]

Ideally, the coding style between forward and backward debug mode (tests too) would match better, but I don't want to make more style only changes now

[yebai]

@penelopeysm, can you help take a look?

[Technici4n]

LGTM

[Technici4n]

Nit: inline?

[sunxd3]

I am following DebugRRule (line 220) here and I think noinline makes more sense to me just from that we want this function to fail and correctly attribute to this definition. Frankly, I am not sure whether the inline hint is doing anything.

[Technici4n]

Sure, either is fine by me

[yebai]

thanks @sunxd3 and @Technici4n -- please feel free to merge and make new releases!