Only seed structurally non-zero entries

Author: devmotion

src/apiutils.jl

@@ -40,48 +40,63 @@ end
     return Expr(:tuple, [:(single_seed(Partials{N,V}, Val{$i}())) for i in 1:N]...)
 end
 
-function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
-               seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
-    duals .= Dual{T,V,N}.(x, Ref(seed))
-    return duals
-end
-
-function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
-               seeds::NTuple{N,Partials{N,V}}) where {T,V,N}
-    dual_inds = 1:N
-    duals[dual_inds] .= Dual{T,V,N}.(view(x,dual_inds), seeds)
-    return duals
-end
-
-# Triangular matrices
-function _nonzero_indices(x::UpperTriangular)
+# Only seed indices that are structurally non-zero
+_structural_nonzero_indices(x::AbstractArray) = eachindex(x)
+function _structural_nonzero_indices(x::UpperTriangular)
     n = size(x, 1)
     return (CartesianIndex(i, j) for j in 1:n for i in 1:j)
 end
-function _nonzero_indices(x::LowerTriangular)
+function _structural_nonzero_indices(x::LowerTriangular)
     n = size(x, 1)
     return (CartesianIndex(i, j) for j in 1:n for i in j:n)
 end
-function seed!(duals::Union{LowerTriangular{Dual{T,V,N}},UpperTriangular{Dual{T,V,N}}}, x, seeds::NTuple{N,Partials{N,V}}) where {T,V,N}
-    for (idx, seed) in zip(_nonzero_indices(duals), seeds)
+_structural_nonzero_indices(x::Diagonal) = diagind(x)
+
+function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
+               seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
+    if eachindex(duals) != eachindex(x)
+        throw(ArgumentError("indices of input array and array of duals are not identical"))
+    end
+    for idx in _structural_nonzero_indices(duals)
         duals[idx] = Dual{T,V,N}(x[idx], seed)
     end
     return duals
 end
 
+function seed!(duals::AbstractArray{Dual{T,V,N}}, x,
+               seeds::NTuple{N,Partials{N,V}}) where {T,V,N}
+    if eachindex(duals) != eachindex(x)
+        throw(ArgumentError("indices of input array and array of duals are not identical"))
+    end
+    for (i, idx) in enumerate(_structural_nonzero_indices(duals))
+        duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
+    end
+    return duals
+end
+
 function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
                seed::Partials{N,V} = zero(Partials{N,V})) where {T,V,N}
+    if eachindex(duals) != eachindex(x)
+        throw(ArgumentError("indices of input array and array of duals are not identical"))
+    end
     offset = index - 1
-    dual_inds = (1:N) .+ offset
-    duals[dual_inds] .= Dual{T,V,N}.(view(x, dual_inds), Ref(seed))
+    idxs = Iterators.drop(_structural_nonzero_indices(duals), offset)
+    for idx in idxs
+        duals[idx] = Dual{T,V,N}(x[idx], seed)
+    end
     return duals
 end
 
 function seed!(duals::AbstractArray{Dual{T,V,N}}, x, index,
                seeds::NTuple{N,Partials{N,V}}, chunksize = N) where {T,V,N}
+    if eachindex(duals) != eachindex(x)
+        throw(ArgumentError("indices of input array and array of duals are not identical"))
+    end
     offset = index - 1
-    seed_inds = 1:chunksize
-    dual_inds = seed_inds .+ offset
-    duals[dual_inds] .= Dual{T,V,N}.(view(x, dual_inds), getindex.(Ref(seeds), seed_inds))
+    idxs = Iterators.drop(_structural_nonzero_indices(duals), offset)
+    for (i, idx) in enumerate(idxs)
+        i > chunksize && break
+        duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
+    end
     return duals
 end

src/gradient.jl

@@ -63,20 +63,20 @@ function extract_gradient!(::Type{T}, result::DiffResult, dual::Dual) where {T}
 end
 
 extract_gradient!(::Type{T}, result::AbstractArray, y::Real) where {T} = fill!(result, zero(y))
-extract_gradient!(::Type{T}, result::AbstractArray, dual::Dual) where {T}= copyto!(result, partials(T, dual))
-
-# Triangular matrices
-function extract_gradient!(::Type{T}, result::Union{UpperTriangular,LowerTriangular}, dual::Dual) where {T}
-    for (idx, p) in zip(_nonzero_indices(result), partials(T, dual))
-        result[idx] = p
+function extract_gradient!(::Type{T}, result::AbstractArray, dual::Dual) where {T}
+    idxs = _structural_nonzero_indices(result)
+    for (i, idx) in enumerate(idxs)
+        result[idx] = partials(T, dual, i)
     end
     return result
 end
 
 function extract_gradient_chunk!(::Type{T}, result, dual, index, chunksize) where {T}
     offset = index - 1
-    for i in 1:chunksize
-        result[i + offset] = partials(T, dual, i)
+    idxs = Iterators.drop(_structural_nonzero_indices(result), offset)
+    for (i, idx) in enumerate(idxs)
+        i > chunksize && break
+        result[idx] = partials(T, dual, i)
     end
     return result
 end

src/prelude.jl

@@ -12,8 +12,15 @@ function Chunk(input_length::Integer, threshold::Integer = DEFAULT_CHUNK_THRESHO
     Base.@nif 12 d->(N == d) d->(Chunk{d}()) d->(Chunk{N}())
 end
 
+_length_structural_nonzero_indices(x::AbstractArray) = length(x)
+function _length_structural_nonzero_indices(x::Union{LowerTriangular,UpperTriangular})
+    n = size(x, 1)
+    return (n * (n + 1)) >> 1
+end
+_length_structural_nonzero_indices(x::Diagonal) = size(x, 1)
+
 function Chunk(x::AbstractArray, threshold::Integer = DEFAULT_CHUNK_THRESHOLD)
-    return Chunk(length(x), threshold)
+    return Chunk(_length_structural_nonzero_indices(x), threshold)
 end
 
 # Constrained to `N <= threshold`, minimize (in order of priority):

test/GradientTest.jl

@@ -227,11 +227,13 @@ end
 end
 
 # issue #738
-@testset "LowerTriangular and UpperTriangular" begin
-    M = rand(3, 3)
-    for T in (LowerTriangular, UpperTriangular)
-        @test ForwardDiff.gradient(sum, T(randn(3, 3))) == T(ones(3, 3))
-        @test ForwardDiff.gradient(x -> dot(M, x), T(randn(3, 3))) == T(M)
+@testset "LowerTriangular, UpperTriangular and Diagonal" begin
+    for n in (3, 10, 20)
+        M = rand(n, n)
+        for T in (LowerTriangular, UpperTriangular, Diagonal)
+            @test ForwardDiff.gradient(sum, T(randn(n, n))) == T(ones(n, n))
+            @test ForwardDiff.gradient(x -> dot(M, x), T(randn(n, n))) == T(M)
+        end
     end
 end