feat: DifferentiateWith + context arguments

[gdalle]

Partial fix for #806, #675.

Since DI doesn't compute derivatives with respect to contexts, in each implementation we need a way to mark the derivative as unknown.

• Enrich DifferentiateWith with a context_wrappers field, denoting how additional arguments beyond x must be taken into account.
• Implement unknown context derivatives in reverse mode:
  • With Zygote by returning a ChainRulesCore.@not_implemented tangent
  • With Mooncake by poisoning the FData and RData of context arguments using NaN
• Implement unknown context derivatives in forward mode:
  • Impossible with ForwardDiff because we don't know right away whether a context argument contains duals or not, and any Constant context containing Duals would make the derivative wrong. We could allow Caches though.
  • Possible with Mooncake but the frule!! doesn't yet exist for DifferentiateWith
• Add DI tests
• Add native tests: impossible because we are missing some tangents?

[gdalle]

@Technici4n I'd love to have some feedback on this, and the related chalk-lab/Mooncake.jl#548 it enables!

[codecov]

Codecov Report

❌ Patch coverage is 95.89041% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 97.27%. Comparing base (f8f29d2) to head (d6645ab).

Files with missing lines	Patch %	Lines
...onInterfaceChainRulesCoreExt/differentiate_with.jl	90.90%	1 Missing ⚠️
...e/ext/DifferentiationInterfaceMooncakeExt/utils.jl	92.30%	1 Missing ⚠️
...rentiationInterface/src/misc/differentiate_with.jl	90.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #975      +/-   ##
==========================================
- Coverage   98.20%   97.27%   -0.93%     
==========================================
  Files         135      131       -4     
  Lines        8000     8008       +8     
==========================================
- Hits         7856     7790      -66     
- Misses        144      218      +74     

Flag	Coverage Δ
DI	97.85% <95.89%> (-1.11%)	⬇️
DIT	95.76% <ø> (-0.47%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Technici4n]

I'm not very familiar with contexts but the basic idea sounds reasonable. Why is the nan poisoning only used now and wasn't previously? Does this PR introduce a new kind of context?

[gdalle]

The idea behind contexts is summed up concisely in https://juliadiff.org/DifferentiationInterface.jl/DifferentiationInterface/stable/explanation/arguments/.
Essentially, DI only computes derivatives with respect to the first argument, but it can take into account additional information, either constants or caches that play a role in the computation but with respect to which the derivatives are not obtained. When writing a rule for a multi-argument function based on an underlying DI call, we thus have to signal that the pullback wrt the first argument is the only one that is actually correct and usable. The rest of the arguments (constants and caches) thus get NaN-ified pullbacks.

[Copilot]

Pull request overview

This PR extends DifferentiateWith to support additional “context” arguments by introducing a context_wrappers field, and updates backend integrations (ChainRulesCore/Zygote, Mooncake) plus internal utilities/tests to preserve array formats when converting between array-of-tuples and tuple-of-arrays.

Changes:

• Extend DifferentiateWith to accept context_wrappers (default ()), update show, and adjust ChainRulesCore/Mooncake/ForwardDiff integrations accordingly.
• Add/adjust internal array conversion utility arroftup_to_tupofarr(tx, x) to preserve the primal’s array type (incl. a GPUArraysCore-specific method).
• Expand test coverage for “wrong-mode” array format preservation and for DifferentiateWith scenarios (normal/Constant/Cache).

Reviewed changes

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

Show a summary per file

File	Description
DifferentiationInterface/src/misc/differentiate_with.jl	Adds context_wrappers to DifferentiateWith, default constructor, updated call/show behavior.
DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/differentiate_with.jl	Updates rrule to support contexts and return unknown tangents for context derivatives.
DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/differentiate_with.jl	Adds Mooncake primitive handling for DifferentiateWith{C} and poisons context derivatives in reverse mode.
DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/utils.jl	Adds nanify* helpers intended to poison Mooncake fdata/rdata for contexts.
DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/differentiate_with.jl	Restricts ForwardDiff override to DifferentiateWith{0} (no contexts).
DifferentiationInterface/src/utils/linalg.jl	Refactors arroftup_to_tupofarr to preserve primal array type via similar(x).
DifferentiationInterface/ext/DifferentiationInterfaceGPUArraysCoreExt/DifferentiationInterfaceGPUArraysCoreExt.jl	Adds GPU specialization for arroftup_to_tupofarr using Adapt.adapt.
DifferentiationInterface/src/first_order/pushforward.jl	Updates calls to arroftup_to_tupofarr to pass the primal output y.
DifferentiationInterface/src/first_order/pullback.jl	Updates calls to arroftup_to_tupofarr to pass the primal input x.
DifferentiationInterface/src/utils/context.jl	Adds DI.call helper used to apply wrappers cleanly in mappings.
DifferentiationInterface/test/Core/Internals/linalg.jl	Adds tests ensuring arroftup_to_tupofarr preserves array/container types.
DifferentiationInterface/test/Core/SimpleFiniteDiff/test.jl	Adds regression tests for array format preservation.
DifferentiationInterface/test/Core/Internals/display.jl	Updates expected DifferentiateWith string representation.
DifferentiationInterface/test/Back/DifferentiateWith/test.jl	Expands DI tests to normal/Constant/Cache scenario variants and updates Mooncake error expectations.
DifferentiationInterface/ext/DifferentiationInterfaceChainRulesCoreExt/DifferentiationInterfaceChainRulesCoreExt.jl	Imports @not_implemented for unknown context tangents.
DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/DifferentiationInterfaceMooncakeExt.jl	Adds Mooncake imports needed for new poisoning utilities.
DifferentiationInterface/Project.toml	Adds Adapt as a weakdep + compat for GPU array conversions.

Comments suppressed due to low confidence (1)

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/differentiate_with.jl:15

• DifferentiateWith supports context arguments, but the ForwardDiff override only handles DifferentiateWith{0}. For DifferentiateWith{C} with C>0, ForwardDiff will fall back to the generic call implementation and will propagate Duals through context arguments if they contain them, which can silently produce incorrect derivatives (the PR description notes this problem). Consider adding a more specific method for DifferentiateWith{C} + Dual that throws an explicit error (or otherwise marks the context derivative as unknown) so ForwardDiff usage with contexts fails loudly instead of computing the wrong result.

function (dw::DI.DifferentiateWith{0})(x::Dual{T, V, N}) where {T, V, N}
    (; f, backend) = dw
    xval = myvalue(T, x)
    tx = mypartials(T, Val(N), x)
    y, ty = DI.value_and_pushforward(f, backend, xval, tx)
    return make_dual(T, y, ty)
end

function (dw::DI.DifferentiateWith{0})(x::AbstractArray{Dual{T, V, N}}) where {T, V, N}
    (; f, backend) = dw
    xval = myvalue(T, x)
    tx = mypartials(T, Val(N), x)
    y, ty = DI.value_and_pushforward(f, backend, xval, tx)
    return make_dual(T, y, ty)
end

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.