From de5bc26d6b8d8ba6b9fe0d4560de0f921e6ad149 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Mon, 7 Oct 2024 12:59:52 +0200
Subject: [PATCH 1/4] Set default `rows_index` and `cols_index` to `nothing`

---
 src/differentiation/compute_jacobian_ad.jl | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl
index f0cae288..1a954cfe 100644
--- a/src/differentiation/compute_jacobian_ad.jl
+++ b/src/differentiation/compute_jacobian_ad.jl
@@ -348,8 +348,8 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
     if FiniteDiff._use_findstructralnz(sparsity)
         rows_index, cols_index = ArrayInterface.findstructralnz(sparsity)
     else
-        rows_index = 1:size(J, 1)
-        cols_index = 1:size(J, 2)
+        rows_index = nothing
+        cols_index = nothing
     end
 
     # fast path if J and sparsity are both AbstractSparseMatrix and have the same sparsity pattern

From ef2c6083eb46cf0051b09cdbfcc87c000b26a653 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Mon, 7 Oct 2024 13:02:35 +0200
Subject: [PATCH 2/4] Better fix

---
 src/differentiation/compute_jacobian_ad.jl | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/differentiation/compute_jacobian_ad.jl b/src/differentiation/compute_jacobian_ad.jl
index 1a954cfe..45148f44 100644
--- a/src/differentiation/compute_jacobian_ad.jl
+++ b/src/differentiation/compute_jacobian_ad.jl
@@ -348,8 +348,9 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
     if FiniteDiff._use_findstructralnz(sparsity)
         rows_index, cols_index = ArrayInterface.findstructralnz(sparsity)
     else
-        rows_index = nothing
-        cols_index = nothing
+        cartind = vec(CartesianIndices(x))
+        rows_index = Base.Iterators.map(first ∘ Tuple, cartind)
+        cols_index = Base.Iterators.map(last ∘ Tuple, cartind)
     end
 
     # fast path if J and sparsity are both AbstractSparseMatrix and have the same sparsity pattern

From ea20b2659235bded268293d02e870fc2066b0660 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Mon, 7 Oct 2024 13:18:34 +0200
Subject: [PATCH 3/4] Same in PolyesterForwardDiff extension

---
 ext/SparseDiffToolsPolyesterExt.jl | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/ext/SparseDiffToolsPolyesterExt.jl b/ext/SparseDiffToolsPolyesterExt.jl
index ea520810..8708646d 100644
--- a/ext/SparseDiffToolsPolyesterExt.jl
+++ b/ext/SparseDiffToolsPolyesterExt.jl
@@ -34,8 +34,9 @@ function polyesterforwarddiff_color_jacobian(J::AbstractMatrix{<:Number}, f::F,
         rows_index = [rows_index[i] for i in 1:length(rows_index)]
         cols_index = [cols_index[i] for i in 1:length(cols_index)]
     else
-        rows_index = 1:nrows
-        cols_index = 1:ncols
+        cartind = vec(CartesianIndices(x))
+        rows_index = Base.Iterators.map(first ∘ Tuple, cartind)
+        cols_index = Base.Iterators.map(last ∘ Tuple, cartind)
     end
 
     if J isa AbstractSparseMatrix
@@ -57,7 +58,7 @@ function polyesterforwarddiff_color_jacobian(J::AbstractMatrix{<:Number}, f::F,
                              if colorvec[cols_index[idx]] == color_i]
                 rows_index_c = rows_index[pick_inds]
                 cols_index_c = cols_index[pick_inds]
-                @inbounds @simd for i in 1:length(rows_index_c)
+                @simd for i in eachindex(rows_index_c, cols_index_c)
                     J[rows_index_c[i], cols_index_c[i]] = dx[rows_index_c[i]]
                 end
                 color_i += 1

From ac17d0ba66758667b2306c78328e94f84a58b950 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Mon, 7 Oct 2024 13:19:48 +0200
Subject: [PATCH 4/4] Less diff

---
 ext/SparseDiffToolsPolyesterExt.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ext/SparseDiffToolsPolyesterExt.jl b/ext/SparseDiffToolsPolyesterExt.jl
index 8708646d..570f6852 100644
--- a/ext/SparseDiffToolsPolyesterExt.jl
+++ b/ext/SparseDiffToolsPolyesterExt.jl
@@ -58,7 +58,7 @@ function polyesterforwarddiff_color_jacobian(J::AbstractMatrix{<:Number}, f::F,
                              if colorvec[cols_index[idx]] == color_i]
                 rows_index_c = rows_index[pick_inds]
                 cols_index_c = cols_index[pick_inds]
-                @simd for i in eachindex(rows_index_c, cols_index_c)
+                @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