Skip to content

Commit a3b145b

Browse files
vectorized partial decompression
add a comment to the very non-obvious vectorized code
1 parent bf44735 commit a3b145b

File tree

1 file changed

+10
-0
lines changed

1 file changed

+10
-0
lines changed

src/differentiation/compute_jacobian_ad.jl

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,7 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
9898
sparsity = jac_cache.sparsity
9999
color_i = 1
100100
chunksize = length(first(first(jac_cache.p)))
101+
fill!(J,zero(eltype(J)))
101102

102103
for i in 1:length(p)
103104
partial_i = p[i]
@@ -107,11 +108,20 @@ function forwarddiff_color_jacobian!(J::AbstractMatrix{<:Number},
107108
rows_index, cols_index, val = findnz(sparsity)
108109
for j in 1:chunksize
109110
dx .= partials.(fx, j)
111+
#=
110112
for k in 1:length(cols_index)
111113
if color[cols_index[k]] == color_i
112114
J[rows_index[k], cols_index[k]] = dx[rows_index[k]]
113115
end
114116
end
117+
118+
or
119+
120+
J[rows_index, cols_index] .+= (color[cols_index] .== color_i) .* dx[rows_index]
121+
122+
+= means requires a zero'd out start
123+
=#
124+
@. setindex!((J,),getindex((J,),rows_index, cols_index) + (getindex((color,),cols_index) == color_i) * getindex((dx,),rows_index),rows_index, cols_index)
115125
color_i += 1
116126
(color_i > maximum(color)) && continue
117127
end

0 commit comments

Comments
 (0)