Make ForwardDiff a weak dependency extension

- Move ForwardDiff-specific code to PreallocationToolsForwardDiffExt
- ForwardDiff is now an optional dependency loaded via package extensions
- Default implementations provided when ForwardDiff is not available
- Uses ref cells to allow extension to modify behavior at runtime
- All ForwardDiff-specific methods for Dual types moved to extension

Author: ChrisRackauckas

Project.toml

@@ -6,14 +6,15 @@ version = "0.4.29"
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 
 [weakdeps]
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5"
 
 [extensions]
+PreallocationToolsForwardDiffExt = "ForwardDiff"
 PreallocationToolsReverseDiffExt = "ReverseDiff"
 PreallocationToolsSparseConnectivityTracerExt = "SparseConnectivityTracer"
 
@@ -43,6 +44,7 @@ julia = "1.10"
 [extras]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
@@ -59,4 +61,4 @@ Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "ADTypes", "Random", "LabelledArrays", "LinearAlgebra", "OrdinaryDiffEq", "Test", "RecursiveArrayTools", "Pkg", "SafeTestsets", "Optimization", "OptimizationOptimJL", "SparseArrays", "Symbolics", "SparseConnectivityTracer"]
+test = ["Aqua", "ADTypes", "ForwardDiff", "Random", "LabelledArrays", "LinearAlgebra", "OrdinaryDiffEq", "Test", "RecursiveArrayTools", "Pkg", "SafeTestsets", "Optimization", "OptimizationOptimJL", "SparseArrays", "Symbolics", "SparseConnectivityTracer"]

ext/PreallocationToolsForwardDiffExt.jl

@@ -0,0 +1,89 @@
+module PreallocationToolsForwardDiffExt
+
+using PreallocationTools
+using ForwardDiff
+using ArrayInterface
+using Adapt
+
+# Initialize on module load
+function __init__()
+    # Set the dual array creator function
+    PreallocationTools.DUAL_ARRAY_CREATOR[] = function(u::AbstractArray{T}, siz,
+            ::Type{Val{chunk_size}}) where {T, chunk_size}
+        ArrayInterface.restructure(u,
+            zeros(ForwardDiff.Dual{Nothing, T, chunk_size},
+                siz...))
+    end
+    
+    # Set the chunk size function to use ForwardDiff's pickchunksize
+    PreallocationTools.CHUNK_SIZE_FUNC[] = ForwardDiff.pickchunksize
+end
+
+# Define chunksize for ForwardDiff.Dual types
+PreallocationTools.chunksize(::Type{ForwardDiff.Dual{T, V, N}}) where {T, V, N} = N
+
+# Define get_tmp methods for ForwardDiff.Dual types
+function PreallocationTools.get_tmp(dc::PreallocationTools.FixedSizeDiffCache, u::T) where {T <: ForwardDiff.Dual}
+    x = reinterpret(T, dc.dual_du)
+    if PreallocationTools.chunksize(T) === PreallocationTools.chunksize(eltype(dc.dual_du))
+        x
+    else
+        @view x[axes(dc.du)...]
+    end
+end
+
+function PreallocationTools.get_tmp(dc::PreallocationTools.FixedSizeDiffCache, u::Type{T}) where {T <: ForwardDiff.Dual}
+    x = reinterpret(T, dc.dual_du)
+    if PreallocationTools.chunksize(T) === PreallocationTools.chunksize(eltype(dc.dual_du))
+        x
+    else
+        @view x[axes(dc.du)...]
+    end
+end
+
+function PreallocationTools.get_tmp(dc::PreallocationTools.FixedSizeDiffCache, u::AbstractArray{T}) where {T <: ForwardDiff.Dual}
+    x = reinterpret(T, dc.dual_du)
+    if PreallocationTools.chunksize(T) === PreallocationTools.chunksize(eltype(dc.dual_du))
+        x
+    else
+        @view x[axes(dc.du)...]
+    end
+end
+
+function PreallocationTools.get_tmp(dc::PreallocationTools.DiffCache, u::T) where {T <: ForwardDiff.Dual}
+    if isbitstype(T)
+        nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du)
+        if nelem > length(dc.dual_du)
+            PreallocationTools.enlargediffcache!(dc, nelem)
+        end
+        PreallocationTools._restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem)))
+    else
+        PreallocationTools._restructure(dc.du, zeros(T, size(dc.du)))
+    end
+end
+
+function PreallocationTools.get_tmp(dc::PreallocationTools.DiffCache, ::Type{T}) where {T <: ForwardDiff.Dual}
+    if isbitstype(T)
+        nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du)
+        if nelem > length(dc.dual_du)
+            PreallocationTools.enlargediffcache!(dc, nelem)
+        end
+        PreallocationTools._restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem)))
+    else
+        PreallocationTools._restructure(dc.du, zeros(T, size(dc.du)))
+    end
+end
+
+function PreallocationTools.get_tmp(dc::PreallocationTools.DiffCache, u::AbstractArray{T}) where {T <: ForwardDiff.Dual}
+    if isbitstype(T)
+        nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du)
+        if nelem > length(dc.dual_du)
+            PreallocationTools.enlargediffcache!(dc, nelem)
+        end
+        PreallocationTools._restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem)))
+    else
+        PreallocationTools._restructure(dc.du, zeros(T, size(dc.du)))
+    end
+end
+
+end

src/PreallocationTools.jl

@@ -1,6 +1,6 @@
 module PreallocationTools
 
-using ForwardDiff, ArrayInterface, Adapt
+using ArrayInterface, Adapt
 using PrecompileTools
 
 struct FixedSizeDiffCache{T <: AbstractArray, S <: AbstractArray}
@@ -9,11 +9,17 @@ struct FixedSizeDiffCache{T <: AbstractArray, S <: AbstractArray}
     any_du::Vector{Any}
 end
 
+# Mutable container to hold dual array creator that can be updated by extension
+const DUAL_ARRAY_CREATOR = Ref{Union{Nothing,Function}}(nothing)
+
 function FixedSizeDiffCache(u::AbstractArray{T}, siz,
         ::Type{Val{chunk_size}}) where {T, chunk_size}
-    x = ArrayInterface.restructure(u,
-        zeros(ForwardDiff.Dual{Nothing, T, chunk_size},
-            siz...))
+    # Try to use ForwardDiff if available, otherwise fallback
+    x = if !isnothing(DUAL_ARRAY_CREATOR[])
+        DUAL_ARRAY_CREATOR[](u, siz, Val{chunk_size})
+    else
+        similar(u, siz...)
+    end
     xany = Any[]
     FixedSizeDiffCache(deepcopy(u), x, xany)
 end
@@ -25,8 +31,18 @@ Builds a `FixedSizeDiffCache` object that stores both a version of the cache for
 and for the `Dual` version of `u`, allowing use of pre-cached vectors with
 forward-mode automatic differentiation.
 """
+# Default chunk size calculation without ForwardDiff
+default_chunk_size(n) = min(n, 12)
+
+# Mutable container to hold chunk size function that can be updated by extension
+const CHUNK_SIZE_FUNC = Ref{Function}(default_chunk_size)
+
+function forwarddiff_compat_chunk_size(n)
+    CHUNK_SIZE_FUNC[](n)
+end
+
 function FixedSizeDiffCache(u::AbstractArray,
-        ::Type{Val{N}} = Val{ForwardDiff.pickchunksize(length(u))}) where {
+        ::Type{Val{N}} = Val{forwarddiff_compat_chunk_size(length(u))}) where {
         N,
 }
     FixedSizeDiffCache(u, size(u), Val{N})
@@ -36,34 +52,10 @@ function FixedSizeDiffCache(u::AbstractArray, N::Integer)
     FixedSizeDiffCache(u, size(u), Val{N})
 end
 
-chunksize(::Type{ForwardDiff.Dual{T, V, N}}) where {T, V, N} = N
+# Generic fallback for chunksize
+chunksize(::Type{T}) where {T} = 0
 
-function get_tmp(dc::FixedSizeDiffCache, u::T) where {T <: ForwardDiff.Dual}
-    x = reinterpret(T, dc.dual_du)
-    if chunksize(T) === chunksize(eltype(dc.dual_du))
-        x
-    else
-        @view x[axes(dc.du)...]
-    end
-end
-
-function get_tmp(dc::FixedSizeDiffCache, u::Type{T}) where {T <: ForwardDiff.Dual}
-    x = reinterpret(T, dc.dual_du)
-    if chunksize(T) === chunksize(eltype(dc.dual_du))
-        x
-    else
-        @view x[axes(dc.du)...]
-    end
-end
-
-function get_tmp(dc::FixedSizeDiffCache, u::AbstractArray{T}) where {T <: ForwardDiff.Dual}
-    x = reinterpret(T, dc.dual_du)
-    if chunksize(T) === chunksize(eltype(dc.dual_du))
-        x
-    else
-        @view x[axes(dc.du)...]
-    end
-end
+# ForwardDiff-specific methods moved to extension
 
 function get_tmp(dc::FixedSizeDiffCache, u::Union{Number, AbstractArray})
     if promote_type(eltype(dc.du), eltype(u)) <: eltype(dc.du)
@@ -103,19 +95,19 @@ function DiffCache(u::AbstractArray{T}, siz, chunk_sizes) where {T}
 end
 
 """
-`DiffCache(u::AbstractArray, N::Int = ForwardDiff.pickchunksize(length(u)); levels::Int = 1)`
+`DiffCache(u::AbstractArray, N::Int = forwarddiff_compat_chunk_size(length(u)); levels::Int = 1)`
 `DiffCache(u::AbstractArray; N::AbstractArray{<:Int})`
 
 Builds a `DiffCache` object that stores both a version of the cache for `u`
 and for the `Dual` version of `u`, allowing use of pre-cached vectors with
 forward-mode automatic differentiation via
-[ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl).
+[ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) (when available).
 Supports nested AD via keyword `levels` or specifying an array of chunk sizes.
 
 The `DiffCache` also supports sparsity detection via
 [SparseConnectivityTracer.jl](https://github.com/adrhill/SparseConnectivityTracer.jl/).
 """
-function DiffCache(u::AbstractArray, N::Int = ForwardDiff.pickchunksize(length(u));
+function DiffCache(u::AbstractArray, N::Int = forwarddiff_compat_chunk_size(length(u));
         levels::Int = 1)
     DiffCache(u, size(u), N * ones(Int, levels))
 end
@@ -133,41 +125,7 @@ const dualcache = DiffCache
 
 Returns the `Dual` or normal cache array stored in `dc` based on the type of `u`.
 """
-function get_tmp(dc::DiffCache, u::T) where {T <: ForwardDiff.Dual}
-    if isbitstype(T)
-        nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du)
-        if nelem > length(dc.dual_du)
-            enlargediffcache!(dc, nelem)
-        end
-        _restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem)))
-    else
-        _restructure(dc.du, zeros(T, size(dc.du)))
-    end
-end
-
-function get_tmp(dc::DiffCache, ::Type{T}) where {T <: ForwardDiff.Dual}
-    if isbitstype(T)
-        nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du)
-        if nelem > length(dc.dual_du)
-            enlargediffcache!(dc, nelem)
-        end
-        _restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem)))
-    else
-        _restructure(dc.du, zeros(T, size(dc.du)))
-    end
-end
-
-function get_tmp(dc::DiffCache, u::AbstractArray{T}) where {T <: ForwardDiff.Dual}
-    if isbitstype(T)
-        nelem = div(sizeof(T), sizeof(eltype(dc.dual_du))) * length(dc.du)
-        if nelem > length(dc.dual_du)
-            enlargediffcache!(dc, nelem)
-        end
-        _restructure(dc.du, reinterpret(T, view(dc.dual_du, 1:nelem)))
-    else
-        _restructure(dc.du, zeros(T, size(dc.du)))
-    end
-end
+# ForwardDiff-specific methods moved to extension
 
 function get_tmp(dc::DiffCache, u::Union{Number, AbstractArray})
     if promote_type(eltype(dc.du), eltype(u)) <: eltype(dc.du)
@@ -291,6 +249,9 @@ Base.getindex(b::GeneralLazyBufferCache, u::T) where {T} = get_tmp(b, u)
 export GeneralLazyBufferCache, FixedSizeDiffCache, DiffCache, LazyBufferCache, dualcache
 export get_tmp
 
+# Export internal functions for extension use (but not public API)
+# These are needed by the ForwardDiff extension
+
 @setup_workload begin
     @compile_workload begin
         # Basic precompilation