Option for Native AD backend APIs to offer direct stability guarantees ?

[AstitvaAggarwal]

Native AD backend APIs often provide better stability and stronger correctness guarantees than DI's layer around them.
eg: chalk-lab/Mooncake.jl#992 - a frule!!/rrule!! inconsistency in nan_tangent_guard causes forward-mode to return NaNwhere reverse-mode returns 0.0. This kind of edge-case divergence between modes is invisible to DI, which wraps each mode independently.

Broadly the concerns are:

• Maintenance burden: each AD backend update can silently break DI's assumptions, requiring downstream fixes (eg Changes in Mooncake 0.5.25 that may affect DI users #986)
• Correctness: bugs like Inconsistent gradients for forward and reverse mode chalk-lab/Mooncake.jl#992 may only surface after multiple layers of wrapping, making diagnosis harder.
• Performance: calling native API is almost always faster as we don't have the DI overhead. (eg slow forward differentiation with Mooncake #854)

Is there a path toward DI where we can either :

• (a) call native backend APIs more directly where they exist, or
• (b) expose a "native-first" mode that bypasses the wrapper for backends that have the latest stable and correct API?

[gdalle]

Yes, the solution is SciML/ADTypes.jl#140. That PR has been asleep for a while following my latest review, feel free to revive it

[gdalle]

But if we're talking specifically about chalk-lab/Mooncake.jl#992, DI doesn't guarantee either that all backends will behave the same at non-differentiable points. Nothing can guarantee that because AD differentiates software implementations, not mathematical functions.

[AstitvaAggarwal]

I agree,that's why maybe we should resort to each separate backends approach to the problem? DI cannot generalize the logic for that problem anyways let the responsibility for that lie with the AD backends.

[gdalle]

If what you need is a way to distinguish between using DI bindings and using whatever the library at hand decides to implement manually for a given backend, the answer is @wsmoses's PR to ADTypes linked above.

[yolhan83]

Hi, just wanted to know if it's possible to just overwrite the method when a native implementation exists for instance it's doing so in Enzyme Forward:

function DI.value_and_gradient(
        f::F,
        prep::EnzymeForwardGradientPrep{SIG, B},
        backend::AutoEnzyme{<:ForwardMode, <:Union{Nothing, Const}},
        x,
        contexts::Vararg{DI.Constant, C},
    ) where {F, SIG, B, C}
    DI.check_prep(f, prep, backend, x, contexts...)
    (; df, context_shadows, basis_shadows) = prep
    mode = forward_withprimal(backend)
    f_and_df = get_f_and_df_prepared!(df, f, backend, Val(B))
    annotated_contexts = translate_prepared!(context_shadows, contexts, Val(B))
    (; derivs, val) = gradient(
        mode, f_and_df, x, annotated_contexts...; chunk = Val(B), shadows = basis_shadows
    )
    return val, first(derivs)
end

Which kinda just call Enzyme.gradient then when an auto diff backend want to speed things up it just need to implement it in their side and overwrite the DI call here but if they don't want to they can stick with the internals.
Meaning it should be done for Enzyme Jacobian, Mooncake value and gradient and soon value and Jacobian.
It does mean that a lot of code will just be dead inside DI just to simplify new backend coming.
Oh I think I just understood AutoDI would mean the dead code should still be testable on backends ?

[gdalle]

I'm not sure I understand what you're asking for @yolhan83, is it related to #634?

[yolhan83]

not really I think, for instance in the MooncakeExt only pullback and pushforward are defined and DI figures out the rest but Mooncake added its own value_and_gradient!! function which may be faster (not sure) or have its own correctness fixes, so can they just overwrite DI.value_and_gradient(!) and DI.gradient(!) in the Extention and everything will work ?

So DI would just be about giving the basic blocks even if its not perfect but it should always be possible to speed things up adding specific cases, so no perf insurence from DI internals.

[gdalle]

Actually DI already uses native APIs whenever they are available and free of bugs (and assuming I've had time to catch up with them, which wasn't easy given the flurry of recent Mooncake changes). Does this docs page and the foldable table there answer your question?

https://juliadiff.org/DifferentiationInterface.jl/DifferentiationInterface/stable/explanation/backends/#Features

EDIT: the gradient API from Mooncake is incorrectly missing from that table, I see what might have misled you. Check out the code in the extension if you're curious, it will always be up-to-date even when I forget to clean up the docs ^^

[yolhan83]

Then where comes the perf diff if its not going through the internals of DI for instance in

┌───────────────────────┬───────┬─────────────┬────────┬───────────────────────────────┬────────────────────────────┬─────────────────────────────────┐
│                       │       │             │        │       t(eval) / t(ref)        │     t(grad) / t(eval)      │        t(grad) / t(ref)         │
│                       │       │             │        │ ─────────┬──────────┬──────── │ ───────┬─────────┬──────── │ ──────────┬───────────┬──────── │
│                 Model │   Dim │  AD Backend │ Linked │     base │  this PR │ speedup │   base │ this PR │ speedup │      base │   this PR │ speedup │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│               Dynamic │    10 │    mooncake │   true │   298.22 │   252.08 │    1.18 │   7.47 │    5.98 │    1.25 │   2226.41 │   1508.19 │    1.48 │
│                   LDA │    12 │ reversediff │   true │  2653.14 │  2006.77 │    1.32 │   2.01 │    2.08 │    0.97 │   5341.90 │   4170.91 │    1.28 │
│   Loop univariate 10k │ 10000 │    mooncake │   true │ 30480.32 │ 27449.35 │    1.11 │   7.43 │    6.78 │    1.09 │ 226363.47 │ 186231.48 │    1.22 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│    Loop univariate 1k │  1000 │    mooncake │   true │  4175.06 │  3612.79 │    1.16 │   5.49 │    5.73 │    0.96 │  22937.72 │  20684.20 │    1.11 │
│      Multivariate 10k │ 10000 │    mooncake │   true │ 35785.74 │ 28379.58 │    1.26 │   8.81 │    8.62 │    1.02 │ 315398.16 │ 244673.66 │    1.29 │
│       Multivariate 1k │  1000 │    mooncake │   true │  3971.96 │  3532.87 │    1.12 │   8.48 │    7.33 │    1.16 │  33689.69 │  25907.11 │    1.30 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│ Simple assume observe │     1 │ forwarddiff │  false │     0.93 │     5.67 │    0.16 │  11.25 │    1.36 │    8.26 │     10.44 │      7.72 │    1.35 │
│           Smorgasbord │   201 │ forwarddiff │  false │   964.64 │   786.01 │    1.23 │  84.86 │   77.19 │    1.10 │  81856.53 │  60673.14 │    1.35 │
│           Smorgasbord │   201 │      enzyme │   true │  1392.77 │  1058.73 │    1.32 │   4.40 │    4.51 │    0.98 │   6128.58 │   4772.68 │    1.28 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│           Smorgasbord │   201 │ forwarddiff │   true │  1342.58 │  1077.43 │    1.25 │  70.09 │   77.68 │    0.90 │  94097.19 │  83693.33 │    1.12 │
│           Smorgasbord │   201 │    mooncake │   true │  1778.52 │  1073.46 │    1.66 │   4.00 │    4.45 │    0.90 │   7113.71 │   4772.68 │    1.49 │
│           Smorgasbord │   201 │ reversediff │   true │  1328.12 │  1087.35 │    1.22 │ 127.27 │  117.72 │    1.08 │ 169031.73 │ 128001.64 │    1.32 │
├───────────────────────┼───────┼─────────────┼────────┼──────────┼──────────┼─────────┼────────┼─────────┼─────────┼───────────┼───────────┼─────────┤
│              Submodel │     1 │    mooncake │   true │     0.92 │     2.47 │    0.37 │  29.26 │    8.56 │    3.42 │     27.04 │     21.13 │    1.28 │
└───────────────────────┴───────┴─────────────┴────────┴──────────┴──────────┴─────────┴────────┴─────────┴─────────┴───────────┴───────────┴─────────┘

that shouldn't happen if those never uses DI power to build the high level APIs from the low ones.

[penelopeysm]

@yolhan83 DynamicPPL maintainer here. and I would point out that those benchmarks are unreliable. Notice how many of the primals (t_eval / t_ref) are also sped up despite the fact that the PR does not touch anything outside of AD. This is just a result of CI flakiness.

In fact, local benchmarking shows that the PR is slower on both modes of Enzyme and reverse Mooncake. Only forward-mode Mooncake has any speedup.

I have already raised this internally.

[yolhan83]

Oh lol sorry then so really no reason of getting out of DI ? and about correctness issues like this should never happen if nothing goes within DI internals; from chalk-lab/Mooncake.jl#992
"""
This appears DifferentiationInterface fallback issue, not an inconsistency in Mooncake’s native sqrt rule.

Mooncake’s native sqrt rules in both forward and reverse mode use nan_tangent_guard, so they intentionally behave
like this at x = 0.0:

seed 1.0 -> Inf
seed 0.0 -> 0.0
Why AutoMooncakeForward() differs

AutoMooncake() uses Mooncake’s native reverse rule for sqrt, so zero cotangent is masked to 0.0.
AutoMooncakeForward() uses DI’s pullback-via-pushforward fallback.
That fallback computes Inf * 0.0, which becomes NaN.
So the inconsistency is a bug or limitation in DI’s forward-backend pullback fallback, not in Mooncake’s native sqrt
rule.
"""

[gdalle]

Of course DI adds some code so the overhead will be non-zero. The real question is whether the overhead is actually measurable in a clean benchmark environment. If you do find such a case and can provide an MWE, I'll always look into it and either fix it to the best of my ability, or tell you why I can't. Thats the standard I have always held myself to, and that's exactly what happened in the slowdown @AstitvaAggarwal pointed to (#854).