Fix color index

Vaibhavdixit02 · Feb 16, 2024 · 9fa3444 · 9fa3444
1 parent 4d0c091
commit 9fa3444
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/ext/SparseDiffToolsPolyesterExt.jl b/ext/SparseDiffToolsPolyesterExt.jl
@@ -45,9 +45,9 @@ function polyesterforwarddiff_color_jacobian(J::AbstractMatrix{<:Number}, f::F,
     end
 
     batch((length(p), min(length(p), Threads.nthreads()))) do _, start, stop
+        color_i = (start - 1) * chunksize + 1
         for i in start:stop
             partial_i = p[i]
-            color_i = i
             t_ = reshape(eltype(t).(vecx, ForwardDiff.Partials.(partial_i)), size(t))
             fx = f(t_)
             for j in 1:chunksize
@@ -60,10 +60,11 @@ function polyesterforwarddiff_color_jacobian(J::AbstractMatrix{<:Number}, f::F,
                 @inbounds @simd for i in 1:length(rows_index_c)
                     J[rows_index_c[i], cols_index_c[i]] = dx[rows_index_c[i]]
                 end
+                color_i += 1
             end
         end
     end
     return J
 end
 
-end
+end
diff --git a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl
@@ -94,7 +94,8 @@ function (alg::ApproximateJacobianSparsity)(
         fx, kwargs...)
 end
 
-function (alg::ApproximateJacobianSparsity)(ad::AutoSparsePolyesterForwardDiff, f::F, fx, x;
+function (alg::ApproximateJacobianSparsity)(
+        ad::AutoSparsePolyesterForwardDiff, f::F, fx, x;
         kwargs...) where {F}
     @unpack ntrials, rng = alg
     ck = __chunksize(ad, x)