Handle Subarrays

Author: avik-pal

ext/SparseDiffToolsEnzymeExt.jl

@@ -1,7 +1,8 @@
 module SparseDiffToolsEnzymeExt
 
 import ArrayInterface: fast_scalar_indexing
-import SparseDiffTools: __f̂, __jacobian!, __gradient, __gradient!, AutoSparseEnzyme
+import SparseDiffTools: __f̂,
+    __maybe_copy_x, __jacobian!, __gradient, __gradient!, AutoSparseEnzyme
 # FIXME: For Enzyme we currently assume reverse mode
 import ADTypes: AutoEnzyme
 using Enzyme
@@ -55,4 +56,6 @@ end
     return J
 end
 
+__maybe_copy_x(::Union{AutoSparseEnzyme, AutoEnzyme}, x::SubArray) = copy(x)
+
 end

src/highlevel/common.jl

@@ -129,3 +129,6 @@ function __init_𝒥(c::AbstractMaybeSparseJacobianCache)
 end
 __init_𝒥(::Nothing, ::Type{T}, fx, x) where {T} = similar(fx, T, length(fx), length(x))
 __init_𝒥(J, ::Type{T}, _, _) where {T} = similar(J, T, size(J, 1), size(J, 2))
+
+__maybe_copy_x(_, x) = x
+__maybe_copy_x(_, ::Nothing) = nothing

src/highlevel/reverse_mode.jl

@@ -34,15 +34,22 @@ function sparse_jacobian!(J::AbstractMatrix, ad, cache::ReverseModeJacobianCache
 end
 
 function __sparse_jacobian_reverse_impl!(J::AbstractMatrix, ad, idx_vec,
-    cache::MatrixColoringResult, f, x, fx = nothing)
+    cache::MatrixColoringResult, f, x)
+    return __sparse_jacobian_reverse_impl!(J, ad, idx_vec, cache, f, nothing, x)
+end
+
+function __sparse_jacobian_reverse_impl!(J::AbstractMatrix, ad, idx_vec,
+    cache::MatrixColoringResult, f, fx, x)
     # If `fx` is `nothing` then assume `f` is not in-place
+    x_ = __maybe_copy_x(ad, x)
+    fx_ = __maybe_copy_x(ad, fx)
     @unpack colorvec, nz_rows, nz_cols = cache
     for c in 1:maximum(colorvec)
         @. idx_vec = colorvec == c
         if fx === nothing
-            gs = __gradient(ad, f, x, idx_vec)
+            gs = __gradient(ad, f, x_, idx_vec)
         else
-            gs = __gradient!(ad, f, fx, x, idx_vec)
+            gs = __gradient!(ad, f, fx_, x_, idx_vec)
         end
         pick_idxs = filter(i -> colorvec[nz_rows[i]] == c, 1:length(nz_rows))
         row_idxs = nz_rows[pick_idxs]