max_norm_perturbation kwarg

[AstitvaAggarwal]

closes #765

This PR adds a max_fd_step keyword argument to test_rule, test_frule_correctness, and test_rrule_correctness. When provided, only finite-difference step sizes ε ≤ max_fd_step are used. This is useful for functions defined on a restricted domain (e.g. log, sqrt, cholesky) where large perturbations would step outside it.

Behaviour change: tangent normalisation in test_frule_correctness

test_frule_correctness now normalises the tangent direction before running finite differences:
ẋ = map(normalize_tangent ∘ tangent, x_ẋ)

Previously the raw tangent from x_ẋ was used directly. With normalisation, ẋ has unit norm, so max_fd_step directly controls the step size in input space (rather than a step scaled by the tangent magnitude). The correctness check is unaffected, the FD estimate and the AD result are both linear in ẋ, so scaling ẋ scales both by the same factor. test_rrule_correctness already normalised its tangent direction; this makes test_frule_correctness consistent with it.

CI Summary — GitHub Actions

Documentation Preview

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

Performance

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 │ 170.0 ns │     1.59 │        1.76 │   0.706 │        3.42 │    6.6 │
│             _sum_1000 │ 972.0 ns │     7.16 │        1.04 │  1990.0 │        40.0 │   1.06 │
│          sum_sin_1000 │  7.25 μs │      3.5 │        1.36 │    1.46 │        10.9 │   1.74 │
│         _sum_sin_1000 │  6.47 μs │      3.5 │        1.77 │   228.0 │        12.4 │   2.05 │
│              kron_sum │ 219.0 μs │     12.7 │         3.6 │    13.4 │       347.0 │   19.7 │
│         kron_view_sum │ 300.0 μs │     11.1 │        5.26 │    24.4 │       313.0 │   11.8 │
│ naive_map_sin_cos_exp │  2.32 μs │     3.08 │        1.54 │ missing │        7.65 │   2.14 │
│       map_sin_cos_exp │  2.29 μs │     3.59 │        1.69 │    1.55 │        6.71 │   2.72 │
│ broadcast_sin_cos_exp │  2.32 μs │     3.19 │        1.59 │    3.81 │        1.45 │   2.12 │
│            simple_mlp │ 352.0 μs │     5.26 │        2.68 │    2.21 │        9.67 │   2.95 │
│                gp_lml │ 172.0 μs │     11.3 │        2.54 │    4.84 │     missing │   5.31 │
│    large_single_block │ 421.0 ns │     6.09 │         1.9 │  4910.0 │        33.3 │   2.05 │
└───────────────────────┴──────────┴──────────┴─────────────┴─────────┴─────────────┴────────┘

[sunxd3]

thanks, I'll take a look today

[sunxd3]

thanks for the effort

[AstitvaAggarwal]

thanks for the reviews @sunxd3! I have tried addressing all of them.

[sunxd3]

for test_frule_correctness now x is normalized, which I think is still correct, but a behavior change, could you write about this change in the PR description?

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!