From 9fa344489dcafac4583ce999439ef61f851ebd9c Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Thu, 15 Feb 2024 16:25:54 -0500
Subject: [PATCH] Fix color index

---
 ext/SparseDiffToolsPolyesterExt.jl            | 5 +++--
 ext/SparseDiffToolsPolyesterForwardDiffExt.jl | 3 ++-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/ext/SparseDiffToolsPolyesterExt.jl b/ext/SparseDiffToolsPolyesterExt.jl
index 9e2a6db4..ea520810 100644
--- 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
\ No newline at end of file
+end
diff --git a/ext/SparseDiffToolsPolyesterForwardDiffExt.jl b/ext/SparseDiffToolsPolyesterForwardDiffExt.jl
index e30c9950..3ea13500 100644
--- 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)