Add resize! dispatch for DiffCache and FixedSizeDiffCache

src/PreallocationTools.jl

@@ -14,7 +14,7 @@ dualarraycreator(args...) = nothing
 
 function FixedSizeDiffCache(u::AbstractArray{T}, siz,
         ::Type{Val{chunk_size}}) where {T, chunk_size}
-   x = dualarraycreator(u, siz, Val{chunk_size})
+    x = dualarraycreator(u, siz, Val{chunk_size})
     xany = Any[]
     FixedSizeDiffCache(deepcopy(u), x, xany)
 end
@@ -233,6 +233,44 @@ function get_tmp(b::GeneralLazyBufferCache, u::T) where {T}
 end
 Base.getindex(b::GeneralLazyBufferCache, u::T) where {T} = get_tmp(b, u)
 
+# resize! methods for PreallocationTools types
+# Note: resize! only works for 1D arrays (vectors)
+function Base.resize!(dc::DiffCache, n::Integer)
+    # Only resize if the array is a vector
+    if dc.du isa AbstractVector
+        resize!(dc.du, n)
+    else
+        throw(ArgumentError("resize! is only supported for DiffCache with vector arrays, got $(typeof(dc.du))"))
+    end
+    # dual_du is often pre-allocated for ForwardDiff dual numbers, 
+    # and may need special handling based on chunk size
+    # Only resize if it's a vector
+    if dc.dual_du isa AbstractVector
+        resize!(dc.dual_du, n)
+    end
+    # Always resize the any_du cache
+    resize!(dc.any_du, n)
+    return dc
+end
+
+function Base.resize!(dc::FixedSizeDiffCache, n::Integer)
+    # Only resize if the array is a vector
+    if dc.du isa AbstractVector
+        resize!(dc.du, n)
+    else
+        throw(ArgumentError("resize! is only supported for FixedSizeDiffCache with vector arrays, got $(typeof(dc.du))"))
+    end
+    # dual_du is often pre-allocated for ForwardDiff dual numbers,
+    # and may need special handling based on chunk size
+    # Only resize if it's a vector
+    if dc.dual_du isa AbstractVector
+        resize!(dc.dual_du, n)
+    end
+    # Always resize the any_du cache
+    resize!(dc.any_du, n)
+    return dc
+end
+
 export GeneralLazyBufferCache, FixedSizeDiffCache, DiffCache, LazyBufferCache, dualcache
 export get_tmp
 

test/core_resizing.jl

@@ -52,3 +52,45 @@ _du = DiffCache(du)
 f = A -> loss(_du, u, A, 0.0)
 analyticalsolution = [3.0 0; 0 0]
 @test ForwardDiff.gradient(f, A) ≈ analyticalsolution
+
+# Test resize! functionality for DiffCache
+@testset "resize! for DiffCache" begin
+    u = rand(10)
+    dc = DiffCache(u)
+
+    # Initial size
+    @test length(dc.du) == 10
+    @test length(dc.any_du) == 0  # Initially empty
+
+    # Resize to larger
+    resize!(dc, 20)
+    @test length(dc.du) == 20
+
+    # Resize to smaller
+    resize!(dc, 5)
+    @test length(dc.du) == 5
+
+    # Test that it returns the cache itself
+    @test resize!(dc, 8) === dc
+end
+
+# Test resize! functionality for FixedSizeDiffCache
+@testset "resize! for FixedSizeDiffCache" begin
+    u = rand(10)
+    dc = FixedSizeDiffCache(u)
+
+    # Initial size
+    @test length(dc.du) == 10
+    @test length(dc.any_du) == 0  # Initially empty
+
+    # Resize to larger
+    resize!(dc, 20)
+    @test length(dc.du) == 20
+
+    # Resize to smaller
+    resize!(dc, 5)
+    @test length(dc.du) == 5
+
+    # Test that it returns the cache itself
+    @test resize!(dc, 8) === dc
+end