relax restriction of x::AbstractVector in threaded_gradient!

JuliaDiff · Apr 12, 2024 · 1c32181 · 1c32181
1 parent a72e010
commit 1c32181
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/src/PolyesterForwardDiff.jl b/src/PolyesterForwardDiff.jl
@@ -49,7 +49,7 @@ function evaluate_chunks!(f::F, (r,Δx,x), start, stop, ::ForwardDiff.Chunk{C},
     end
 end
 
-function threaded_gradient!(f::F, Δx::AbstractVector, x::AbstractVector, ::ForwardDiff.Chunk{C}, check = Val{false}()) where {F,C}
+function threaded_gradient!(f::F, Δx::AbstractArray, x::AbstractArray, ::ForwardDiff.Chunk{C}, check = Val{false}()) where {F,C}
     N = length(x)
     d = cld_fast(N, C)
     r = Ref{eltype(Δx)}()

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -37,3 +37,12 @@ ForwardDiff.jacobian!(dxref, g!, yref, x, ForwardDiff.JacobianConfig(g!, yref, x
 PolyesterForwardDiff.threaded_jacobian!(g!, y, dx, x, ForwardDiff.Chunk(8),Val{true}());
 @test dx ≈ dxref
 @test y ≈ yref
+
+
+X = randn(10,80);
+dXref = similar(x);
+dX = similar(x);
+ForwardDiff.gradient!(dXref, f, X, ForwardDiff.GradientConfig(f, x, ForwardDiff.Chunk(8), nothing));
+PolyesterForwardDiff.threaded_gradient!(f, dX, X, ForwardDiff.Chunk(8));
+
+@test dX ≈ dXref