Merge pull request #242 from vpuri3/resize

overload Base.resize!

Author: ChrisRackauckas

src/differentiation/jaches_products.jl

@@ -228,6 +228,16 @@ function (L::FwdModeAutoDiffVecProd)(dv, v, p, t)
     L.vecprod!(dv, L.f, L.u, v, L.cache...)
 end
 
+function Base.resize!(L::FwdModeAutoDiffVecProd, n::Integer)
+
+    static_hasmethod(resize!, typeof((L.f, n))) && resize!(L.f, n)
+    resize!(L.u, n)
+
+    for v in L.cache
+        resize!(v, n)
+    end
+end
+
 function JacVec(f, u::AbstractArray, p = nothing, t = nothing;
                 autodiff = AutoForwardDiff(), tag = DeivVecTag(), kwargs...)
     cache, vecprod, vecprod! = if autodiff isa AutoFiniteDiff

src/differentiation/vecjac_products.jl

@@ -87,6 +87,16 @@ function (L::RevModeAutoDiffVecProd{ad, true, false})(dv, v, p, t) where {ad}
     L.vecprod!(dv, L.f, L.u, v, L.cache...)
 end
 
+function Base.resize!(L::RevModeAutoDiffVecProd, n::Integer)
+
+    static_hasmethod(resize!, typeof((L.f, n))) && resize!(L.f, n)
+    resize!(L.u, n)
+
+    for v in L.cache
+        resize!(v, n)
+    end
+end
+
 function VecJac(f, u::AbstractArray, p = nothing, t = nothing; autodiff = AutoFiniteDiff(),
                 kwargs...)
     vecprod, vecprod! = if autodiff isa AutoFiniteDiff

test/test_jaches_products.jl

@@ -32,6 +32,9 @@ function _h(dy, x)
     FiniteDiff.finite_difference_gradient!(dy, _g, x)
 end
 
+f2(x) = 2x
+f2(y, x) = (copy!(y, x); lmul!(2, y); y)
+
 # Make functions state-dependent for operator tests 
 
 include("update_coeffs_testutils.jl")
@@ -138,6 +141,17 @@ L = JacVec(f, copy(x), 1.0, 1.0)
 L = JacVec(f, copy(x), 1.0, 1.0; tag = MyTag())
 @test get_tag(L.op.cache[1]) === ForwardDiff.Tag{MyTag, eltype(x)}
 
+# Resize test
+for M in (100, 400)
+    L = JacVec(f2, copy(x), 1.0, 1.0)
+    resize!(L, M)
+    _x = resize!(copy(x), M)
+    _u = rand(M)
+
+    @test L * _u ≈ auto_jacvec(f2, _x, _u)
+    _v = zeros(M); @test mul!(_v, L, _u) ≈ auto_jacvec(f2, _x, _u)
+end
+
 @info "HesVec"
 
 L = HesVec(g, copy(x), 1.0, 1.0, autodiff = AutoFiniteDiff())

test/test_vecjac_products.jl

@@ -63,4 +63,21 @@ actual_jac = Zygote.jacobian(f, v)[1]
 # Test that x and v were not mutated
 @test x ≈ x0
 @test v ≈ v0
+
+# Resize test
+f2(x) = 2x
+f2(y, x) = (copy!(y, x); lmul!(2, y); y)
+
+for M in (100, 400)
+    L = VecJac(f2, copy(x), 1.0f0, 1.0f0; autodiff = AutoZygote())
+    resize!(L, M)
+    _x = resize!(copy(x), M)
+    _u = rand(M)
+
+    J2 = Zygote.jacobian(f2, _x)[1]
+
+    @test L * _u ≈ J2' * _u rtol=1e-6
+    _v = zeros(M); @test mul!(_v, L, _u) ≈ J2' * _u rtol=1e-6
+end
+
 #