Apply suggestions from code review

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: gdalle

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/differentiate_with.jl

@@ -1,10 +1,12 @@
 const NumberOrArray = Union{Number, AbstractArray{<:Number}}
-@is_primitive MinimalCtx Tuple{DI.DifferentiateWith{0}, Any}
-@is_primitive MinimalCtx Tuple{DI.DifferentiateWith{1}, Any, Any}
-@is_primitive MinimalCtx Tuple{DI.DifferentiateWith{2}, Any, Any, Any}
-@is_primitive MinimalCtx Tuple{DI.DifferentiateWith{3}, Any, Any, Any, Any}
-# TODO: generate more cases programmatically
 
+# Mark DifferentiateWith with a range of context arities as primitives.
+# For C contexts, the corresponding call tuple type is
+# Tuple{DI.DifferentiateWith{C}, Any, Vararg{Any, C}}:
+# one slot for the primal input x and C slots for contexts.
+for C in 0:16
+    @eval @is_primitive MinimalCtx Tuple{DI.DifferentiateWith{$C}, Vararg{Any, $(C + 1)}}
+end
 struct MooncakeDifferentiateWithError <: Exception
     F::Type
     X::Type
@@ -37,7 +39,7 @@ function Mooncake.rrule!!(
     # output is a vector, so we need to use the vector pullback
     function pullback_array!!(dy::NoRData)
         dx = DI.pullback(f, backend, primal_x, (y.dx,), wrapped_primal_contexts...) |> only
-        @assert rdata(only(dx)) isa rdata_type(tangent_type(typeof(primal_x)))
+        @assert rdata(dx) isa rdata_type(tangent_type(typeof(primal_x)))
         rc = nanify_fdata_and_rdata!!(contexts...)
         return (NoRData(), rdata(dx), rc...)
     end

DifferentiationInterface/ext/DifferentiationInterfaceMooncakeExt/utils.jl

@@ -19,8 +19,9 @@ function zero_tangent_or_primal(x, backend::AnyAutoMooncake)
 end
 
 nanify(x::AbstractFloat) = convert(typeof(x), NaN)
-nanify(x::AbstractArray) = map(nan_tangent, x)
-nanify(x::Union{Tuple, NamedTuple}) = map(nan_tangent, x)
+nanify(x::AbstractArray) = map(nanify, x)
+nanify(x::NamedTuple) = NamedTuple{keys(x)}(map(nanify, values(x)))
+nanify(x::Tuple) = map(nanify, x)
 nanify(::NoFData) = NoFData()
 nanify(::NoRData) = NoRData()
 

DifferentiationInterface/src/misc/differentiate_with.jl

@@ -76,6 +76,29 @@ struct DifferentiateWith{C, F, B <: AbstractADType, N <: NTuple{C, Any}}
     context_wrappers::N
 end
 
+function DifferentiateWith(
+    f::F,
+    backend::B,
+    context_wrappers::NTuple{C, Any},
+) where {F, B <: AbstractADType, C}
+    for (i, wrapper) in pairs(context_wrappers)
+        # Accept typical constructor-like values: functions or types.
+        if !(wrapper isa Function || wrapper isa Type)
+            throw(
+                ArgumentError(
+                    "Each context wrapper must be a callable object or type " *
+                    "(e.g., a wrapper constructor like `Constant` or `Cache`), " *
+                    "but element $i has type $(typeof(wrapper)).",
+                ),
+            )
+        end
+    end
+    return DifferentiateWith{C, F, B, typeof(context_wrappers)}(
+        f,
+        backend,
+        context_wrappers,
+    )
+end
 DifferentiateWith(f::F, backend::AbstractADType) where {F} = DifferentiateWith(f, backend, ())
 
 function (dw::DifferentiateWith{C})(x, args::Vararg{Any, C}) where {C}