Merge pull request #282 from avik-pal/ap/test

Fix master

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "SparseDiffTools"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 authors = ["Pankaj Mishra <[email protected]>", "Chris Rackauckas <[email protected]>"]
-version = "2.15.0"
+version = "2.15.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

ext/SparseDiffToolsZygoteExt.jl

@@ -45,7 +45,8 @@ end
 
 ### Jac, Hes products
 
-function numback_hesvec!(dy, f::F, x, v, cache1 = similar(v), cache2 = similar(v), cache3 = similar(v)) where {F}
+function numback_hesvec!(dy, f::F, x, v, cache1 = similar(v), cache2 = similar(v),
+        cache3 = similar(v)) where {F}
     g = let f = f
         (dx, x) -> dx .= first(Zygote.gradient(f, x))
     end

src/SparseDiffTools.jl

@@ -14,7 +14,7 @@ import ADTypes: AbstractADType, AutoSparseZygote, AbstractSparseForwardMode,
 import ForwardDiff: Dual, jacobian, partials, DEFAULT_CHUNK_THRESHOLD
 # Array Packages
 using ArrayInterface, SparseArrays
-import ArrayInterface: matrix_colors
+import ArrayInterface: matrix_colors, allowed_setindex!
 import StaticArrays
 import StaticArrays: StaticArray, SArray, MArray, Size
 # Others

src/differentiation/jaches_products.jl

@@ -32,8 +32,8 @@ function auto_jacvec(f, x, v)
     vec(partials.(vec(f(y)), 1))
 end
 
-function num_jacvec!(dy, f, x, v, cache1 = similar(v), cache2 = similar(v), cache3 = similar(v);
-        compute_f0 = true)
+function num_jacvec!(dy, f, x, v, cache1 = similar(v), cache2 = similar(v),
+        cache3 = similar(v); compute_f0 = true)
     vv = reshape(v, axes(x))
     compute_f0 && (f(cache1, x))
     T = eltype(x)
@@ -134,7 +134,8 @@ function autonum_hesvec(f, x, v)
     partials.(g(Dual{DeivVecTag}.(x, v)), 1)
 end
 
-function num_hesvecgrad!(dy, g, x, v, cache1 = similar(v), cache2 = similar(v), cache3 = similar(v))
+function num_hesvecgrad!(dy, g, x, v, cache1 = similar(v), cache2 = similar(v),
+        cache3 = similar(v))
     T = eltype(x)
     # Should it be min? max? mean?
     ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))

src/differentiation/vecjac_products.jl

@@ -1,5 +1,5 @@
-function num_vecjac!(du, f::F, x, v, cache1 = similar(v), cache2 = similar(v), cache3 = similar(v);
-        compute_f0 = true) where {F}
+function num_vecjac!(du, f::F, x, v, cache1 = similar(v), cache2 = similar(v),
+        cache3 = similar(x); compute_f0 = true) where {F}
     compute_f0 && (f(cache1, x))
     T = eltype(x)
     # Should it be min? max? mean?
@@ -15,19 +15,35 @@ function num_vecjac!(du, f::F, x, v, cache1 = similar(v), cache2 = similar(v), c
     return du
 end
 
+# Special Non-Allocating case for StaticArrays
+function num_vecjac(f::F, x::SArray, v::SArray, f0 = nothing) where {F}
+    f0 === nothing ? (_f0 = f(x)) : (_f0 = f0)
+    vv = reshape(v, axes(_f0))
+    T = eltype(x)
+    ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
+    du = zeros(typeof(x))
+    for i in 1:length(x)
+        cache = Base.setindex(x, x[i] + ϵ, i)
+        f0 = f(cache)
+        du = Base.setindex(du, (((f0 .- _f0) ./ ϵ)' * vv), i)
+    end
+    return du
+end
+
 function num_vecjac(f::F, x, v, f0 = nothing) where {F}
     f0 === nothing ? (_f0 = f(x)) : (_f0 = f0)
     vv = reshape(v, axes(_f0))
     T = eltype(x)
     # Should it be min? max? mean?
     ϵ = sqrt(eps(real(T))) * max(one(real(T)), abs(norm(x)))
     du = similar(x)
-    cache = copy(x)
+    cache = similar(x)
+    copyto!(cache, x)
     for i in 1:length(x)
-        cache[i] += ϵ
-        f0 = f(x)
-        cache[i] = x[i]
-        du[i] = (((f0 .- _f0) ./ ϵ)' * vv)[1]
+        cache = allowed_setindex!(cache, x[i] + ϵ, i)
+        f0 = f(cache)
+        cache = allowed_setindex!(cache, x[i], i)
+        du = allowed_setindex!(du, (((f0 .- _f0) ./ ϵ)' * vv)[1], i)
     end
     return vec(du)
 end
@@ -93,7 +109,7 @@ function VecJac(f, u::AbstractArray, p = nothing, t = nothing; fu = nothing,
 end
 
 function _vecjac(f::F, fu, u, autodiff::AutoFiniteDiff) where {F}
-    cache = (similar(fu), similar(fu), similar(fu))
+    cache = (similar(fu), similar(fu), similar(u))
     pullback = nothing
     return AutoDiffVJP(f, u, cache, autodiff, pullback)
 end

test/test_vecjac_products.jl

@@ -1,5 +1,6 @@
-using SparseDiffTools, Zygote
+using SparseDiffTools, Zygote, ForwardDiff
 using LinearAlgebra, Test
+using StaticArrays
 
 using Random
 Random.seed!(123)
@@ -170,3 +171,17 @@ L = VecJac(f3_iip, copy(x); autodiff = AutoFiniteDiff(), fu = copy(y))
 L = VecJac(f3_oop, copy(x); autodiff = AutoZygote())
 @test size(L) == (length(x), length(y))
 @test L * y ≈ Zygote.jacobian(f3_oop, copy(x))[1]' * y
+
+@info "Testing StaticArrays"
+
+const A_sa = rand(SMatrix{4, 4, Float32})
+_f_sa(x) = A_sa * (x .^ 2)
+
+u = rand(SVector{4, Float32})
+v = rand(SVector{4, Float32})
+
+J = ForwardDiff.jacobian(_f_sa, u)
+Jᵀv_true = J' * v
+
+@test num_vecjac(_f_sa, u, v) isa SArray
+@test num_vecjac(_f_sa, u, v)≈Jᵀv_true atol=1e-3