Fix tests

Author: devmotion

src/apiutils.jl

@@ -41,46 +41,55 @@ end
 end
 
 # Only seed indices that are structurally non-zero
-_structural_nonzero_indices(x::AbstractArray) = eachindex(x)
-function _structural_nonzero_indices(x::UpperTriangular)
+structural_eachindex(x::AbstractArray) = structural_eachindex(x, x)
+function structural_eachindex(x::AbstractArray, y::AbstractArray)
+    require_one_based_indexing(x, y)
+    eachindex(x, y)
+end
+function structural_eachindex(x::UpperTriangular, y::AbstractArray)
+    require_one_based_indexing(x, y)
+    if size(x) != size(y)
+        throw(DimensionMismatch())
+    end
     n = size(x, 1)
     return (CartesianIndex(i, j) for j in 1:n for i in 1:j)
 end
-function _structural_nonzero_indices(x::LowerTriangular)
+function structural_eachindex(x::LowerTriangular, y::AbstractArray)
+    require_one_based_indexing(x, y)
+    if size(x) != size(y)
+        throw(DimensionMismatch())
+    end
     n = size(x, 1)
     return (CartesianIndex(i, j) for j in 1:n for i in j:n)
 end
-_structural_nonzero_indices(x::Diagonal) = diagind(x)
+function structural_eachindex(x::Diagonal, y::AbstractArray)
+    require_one_based_indexing(x, y)
+    if size(x) != size(y)
+        throw(DimensionMismatch())
+    end
+    return diagind(x)
+end
 
 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)
+    for idx in structural_eachindex(duals, x)
         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))
+    for (i, idx) in zip(1:N, structural_eachindex(duals, x))
         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
-    idxs = Iterators.drop(_structural_nonzero_indices(duals), offset)
+    idxs = Iterators.drop(structural_eachindex(duals, x), offset)
     for idx in idxs
         duals[idx] = Dual{T,V,N}(x[idx], seed)
     end
@@ -89,13 +98,9 @@ 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
-    idxs = Iterators.drop(_structural_nonzero_indices(duals), offset)
-    for (i, idx) in enumerate(idxs)
-        i > chunksize && break
+    idxs = Iterators.drop(structural_eachindex(duals, x), offset)
+    for (i, idx) in zip(1:chunksize, idxs)
         duals[idx] = Dual{T,V,N}(x[idx], seeds[i])
     end
     return duals

src/gradient.jl

@@ -64,18 +64,17 @@ end
 
 extract_gradient!(::Type{T}, result::AbstractArray, y::Real) where {T} = fill!(result, zero(y))
 function extract_gradient!(::Type{T}, result::AbstractArray, dual::Dual) where {T}
-    idxs = _structural_nonzero_indices(result)
-    for (i, idx) in enumerate(idxs)
+    idxs = structural_eachindex(result)
+    for (i, idx) in zip(1:npartials(dual), 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
-    idxs = Iterators.drop(_structural_nonzero_indices(result), offset)
-    for (i, idx) in enumerate(idxs)
-        i > chunksize && break
+    idxs = Iterators.drop(structural_eachindex(result), offset)
+    for (i, idx) in zip(1:chunksize, idxs)
         result[idx] = partials(T, dual, i)
     end
     return result