Make mutable pre-allocated Jacobian default oop autodiff method whenever J is mutable.

Author: chenwilliam77

src/differentiation/compute_jacobian_ad.jl

@@ -68,13 +68,16 @@ end
                 jac_prototype = nothing,
                 chunksize = nothing,
                 dx = sparsity === nothing && jac_prototype === nothing ? nothing : copy(x)) #if dx is nothing, we will estimate dx at the cost of a function call
+    @show typeof(x)
+
     if sparsity === nothing && jac_prototype === nothing || !ArrayInterface.ismutable(x)
         cfg = chunksize === nothing ? ForwardDiff.JacobianConfig(f, x) : ForwardDiff.JacobianConfig(f, x, ForwardDiff.Chunk(getsize(chunksize)))
         return ForwardDiff.jacobian(f, x, cfg)
     end
     if dx isa Nothing
         dx = f(x)
     end
+    @show "Line 80"
     forwarddiff_color_jacobian(f,x,ForwardColorJacCache(f,x,chunksize,dx=dx,colorvec=colorvec,sparsity=sparsity),jac_prototype)
 end
 
@@ -84,11 +87,12 @@ end
                 sparsity = nothing,
                 jac_prototype = nothing,
                 chunksize = nothing,
-                dx = similar(x, size(J, 1))) #if dx is nothing, we will estimate dx at the cost of a function call
+                dx = similar(x, size(J, 1))) #dx kwarg can be used to avoid re-allocating dx every time
     if sparsity === nothing && jac_prototype === nothing || !ArrayInterface.ismutable(x)
         cfg = chunksize === nothing ? ForwardDiff.JacobianConfig(f, x) : ForwardDiff.JacobianConfig(f, x, ForwardDiff.Chunk(getsize(chunksize)))
         return ForwardDiff.jacobian(f, x, cfg)
     end
+    @show "Line 95"
     forwarddiff_color_jacobian(J,f,x,ForwardColorJacCache(f,x,chunksize,dx=dx,colorvec=colorvec,sparsity=sparsity),jac_prototype)
 end
 
@@ -99,10 +103,18 @@ function forwarddiff_color_jacobian(f,x::AbstractArray{<:Number},jac_cache::Forw
 
     J = jac_prototype isa Nothing ? (sparsity isa Nothing ? false .* vec(dx) .* vecx' : zeros(eltype(x),size(sparsity))) : zero(jac_prototype)
 
-    forwarddiff_color_jacobian(J, f, x, jac_cache, jac_prototype)
+    @show typeof(J)
+    if ArrayInterface.ismutable(J) # Whenever J is mutable, we mutate it to avoid allocations
+        @show "Line 108"
+        forwarddiff_color_jacobian(J, f, x, jac_cache, jac_prototype)
+    else
+        @show "Line 111"
+        forwarddiff_color_jacobian_immutable(J, f, x, jac_cache, jac_prototype)
+    end
 end
 
-function forwarddiff_color_jacobian(J::AbstractArray{<:Number},f,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache,jac_prototype=nothing)
+# When J is mutable, this version of forwarddiff_color_jacobian will mutate J to avoid allocations
+function forwarddiff_color_jacobian(J::AbstractMatrix{<:Number},f,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache)
     t = jac_cache.t
     dx = jac_cache.dx
     p = jac_cache.p
@@ -141,7 +153,11 @@ function forwarddiff_color_jacobian(J::AbstractArray{<:Number},f,x::AbstractArra
                     cols_index_c = vcat(cols_index_c,zeros(Int,nrows-len_rows))[perm_rows]
                     Ji = [j==cols_index_c[i] ? dx[i] : false for i in 1:nrows, j in 1:ncols]
                 end
-                J = J + Ji
+                if j == 1 && i == 1
+                    J .= Ji # overwrite pre-allocated matrix
+                else
+                    J .+= Ji
+                end
                 color_i += 1
                 (color_i > maxcolor) && return J
             end
@@ -150,14 +166,19 @@ function forwarddiff_color_jacobian(J::AbstractArray{<:Number},f,x::AbstractArra
                 col_index = (i-1)*chunksize + j
                 (col_index > ncols) && return J
                 Ji = mapreduce(i -> i==col_index ? partials.(vec(fx), j) : adapt(parameterless_type(J),zeros(eltype(J),nrows)), hcat, 1:ncols)
-                J = J + (size(Ji)!=size(J) ? reshape(Ji,size(J)) : Ji) #branch when size(dx) == (1,) => size(Ji) == (1,) while size(J) == (1,1)
+                if j == 1 && i == 1
+                    J .= (size(Ji)!=size(J) ? reshape(Ji,size(J)) : Ji) # overwrite pre-allocated matrix
+                else
+                    J .+= (size(Ji)!=size(J) ? reshape(Ji,size(J)) : Ji) #branch when size(dx) == (1,) => size(Ji) == (1,) while size(J) == (1,1)
+                end
             end
         end
     end
     J
 end
 
-function forwarddiff_color_jacobian(J::SparseMatrixCSC{<:Number},f,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache,jac_prototype=nothing)
+# When J is immutable, this version of forwarddiff_color_jacobian will avoid mutating J
+function forwarddiff_color_jacobian_immutable(J::AbstractArray{<:Number},f,x::AbstractArray{<:Number},jac_cache::ForwardColorJacCache)
     t = jac_cache.t
     dx = jac_cache.dx
     p = jac_cache.p
@@ -187,13 +208,16 @@ function forwarddiff_color_jacobian(J::SparseMatrixCSC{<:Number},f,x::AbstractAr
                 pick_inds = [i for i in 1:length(rows_index) if colorvec[cols_index[i]] == color_i]
                 rows_index_c = rows_index[pick_inds]
                 cols_index_c = cols_index[pick_inds]
-                Ji = sparse(rows_index_c, cols_index_c, dx[rows_index_c],nrows,ncols)
-                # J = J + Ji
-                if j == 1 && i == 1
-                    J .= Ji # overwrite pre-allocated matrix
+                if J isa SparseMatrixCSC
+                    Ji = sparse(rows_index_c, cols_index_c, dx[rows_index_c],nrows,ncols)
                 else
-                    J .+= Ji
+                    len_rows = length(pick_inds)
+                    unused_rows = setdiff(1:nrows,rows_index_c)
+                    perm_rows = sortperm(vcat(rows_index_c,unused_rows))
+                    cols_index_c = vcat(cols_index_c,zeros(Int,nrows-len_rows))[perm_rows]
+                    Ji = [j==cols_index_c[i] ? dx[i] : false for i in 1:nrows, j in 1:ncols]
                 end
+                J = J + Ji
                 color_i += 1
                 (color_i > maxcolor) && return J
             end
@@ -202,12 +226,7 @@ function forwarddiff_color_jacobian(J::SparseMatrixCSC{<:Number},f,x::AbstractAr
                 col_index = (i-1)*chunksize + j
                 (col_index > ncols) && return J
                 Ji = mapreduce(i -> i==col_index ? partials.(vec(fx), j) : adapt(parameterless_type(J),zeros(eltype(J),nrows)), hcat, 1:ncols)
-                # J = J + (size(Ji)!=size(J) ? reshape(Ji,size(J)) : Ji) #branch when size(dx) == (1,) => size(Ji) == (1,) while size(J) == (1,1)
-                if j == 1 && i == 1
-                    J .= (size(Ji)!=size(J) ? reshape(Ji,size(J)) : Ji) # overwrite pre-allocated matrix
-                else
-                    J .+= (size(Ji)!=size(J) ? reshape(Ji,size(J)) : Ji) #branch when size(dx) == (1,) => size(Ji) == (1,) while size(J) == (1,1)
-                end
+                J = J + (size(Ji)!=size(J) ? reshape(Ji,size(J)) : Ji) #branch when size(dx) == (1,) => size(Ji) == (1,) while size(J) == (1,1)
             end
         end
     end