Merge pull request #23 from longemen3000/patch-2

Suppor more array shapes in `threaded_gradient!`

Author: ChrisRackauckas

README.md

@@ -3,7 +3,6 @@
 [![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://JuliaSIMD.github.io/Polyester.jl/stable)
 [![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://JuliaSIMD.github.io/Polyester.jl/dev)
 [![CI](https://github.com/JuliaDiff/PolyesterForwardDiff.jl/actions/workflows/CI.yml/badge.svg)](https://github.com/JuliaDiff/PolyesterForwardDiff.jl/actions/workflows/CI.yml)
-[![CI-Nightly](https://github.com/JuliaDiff/PolyesterForwardDiff.jl/actions/workflows/CI-julia-nightly.yml/badge.svg)](https://github.com/JuliaDiff/PolyesterForwardDiff.jl/actions/workflows/CI-julia-nightly.yml)
 [![Coverage](https://codecov.io/gh/JuliaDiff/PolyesterForwardDiff.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/JuliaDiff/PolyesterForwardDiff.jl)
 
 

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)}()

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