Base vector vs chunk mode on structural length

Author: devmotion

src/gradient.jl

@@ -16,7 +16,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba
 function gradient(f::F, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f, x), ::Val{CHK}=Val{true}()) where {F, T, CHK}
     require_one_based_indexing(x)
     CHK && checktag(T, f, x)
-    if chunksize(cfg) == length(x)
+    if chunksize(cfg) == structural_length(x)
         return vector_mode_gradient(f, x, cfg)
     else
         return chunk_mode_gradient(f, x, cfg)
@@ -35,7 +35,7 @@ This method assumes that `isa(f(x), Real)`.
 function gradient!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArray, cfg::GradientConfig{T} = GradientConfig(f, x), ::Val{CHK}=Val{true}()) where {T, CHK, F}
     result isa DiffResult ? require_one_based_indexing(x) : require_one_based_indexing(result, x)
     CHK && checktag(T, f, x)
-    if chunksize(cfg) == length(x)
+    if chunksize(cfg) == structural_length(x)
         vector_mode_gradient!(result, f, x, cfg)
     else
         chunk_mode_gradient!(result, f, x, cfg)
@@ -114,10 +114,10 @@ end
 
 function chunk_mode_gradient_expr(result_definition::Expr)
     return quote
-        @assert length(x) >= N "chunk size cannot be greater than length(x) ($(N) > $(length(x)))"
+        @assert structural_length(x) >= N "chunk size cannot be greater than ForwardDiff.structural_length(x) ($(N) > $(structural_length(x)))"
 
         # precalculate loop bounds
-        xlen = length(x)
+        xlen = structural_length(x)
         remainder = xlen % N
         lastchunksize = ifelse(remainder == 0, N, remainder)
         lastchunkindex = xlen - lastchunksize + 1

src/jacobian.jl

@@ -18,7 +18,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba
 function jacobian(f::F, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T,CHK}
     require_one_based_indexing(x)
     CHK && checktag(T, f, x)
-    if chunksize(cfg) == length(x)
+    if chunksize(cfg) == structural_length(x)
         return vector_mode_jacobian(f, x, cfg)
     else
         return chunk_mode_jacobian(f, x, cfg)
@@ -36,7 +36,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba
 function jacobian(f!::F, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {F,T, CHK}
     require_one_based_indexing(y, x)
     CHK && checktag(T, f!, x)
-    if chunksize(cfg) == length(x)
+    if chunksize(cfg) == structural_length(x)
         return vector_mode_jacobian(f!, y, x, cfg)
     else
         return chunk_mode_jacobian(f!, y, x, cfg)
@@ -57,7 +57,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba
 function jacobian!(result::Union{AbstractArray,DiffResult}, f::F, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f, x), ::Val{CHK}=Val{true}()) where {F,T, CHK}
     result isa DiffResult ? require_one_based_indexing(x) : require_one_based_indexing(result, x)
     CHK && checktag(T, f, x)
-    if chunksize(cfg) == length(x)
+    if chunksize(cfg) == structural_length(x)
         vector_mode_jacobian!(result, f, x, cfg)
     else
         chunk_mode_jacobian!(result, f, x, cfg)
@@ -78,7 +78,7 @@ Set `check` to `Val{false}()` to disable tag checking. This can lead to perturba
 function jacobian!(result::Union{AbstractArray,DiffResult}, f!::F, y::AbstractArray, x::AbstractArray, cfg::JacobianConfig{T} = JacobianConfig(f!, y, x), ::Val{CHK}=Val{true}()) where {F,T,CHK}
     result isa DiffResult ? require_one_based_indexing(y, x) : require_one_based_indexing(result, y, x)
     CHK && checktag(T, f!, x)
-    if chunksize(cfg) == length(x)
+    if chunksize(cfg) == structural_length(x)
         vector_mode_jacobian!(result, f!, y, x, cfg)
     else
         chunk_mode_jacobian!(result, f!, y, x, cfg)
@@ -169,10 +169,10 @@ const JACOBIAN_ERROR = DimensionMismatch("jacobian(f, x) expects that f(x) is an
 function jacobian_chunk_mode_expr(work_array_definition::Expr, compute_ydual::Expr,
                                   result_definition::Expr, y_definition::Expr)
     return quote
-        @assert length(x) >= N "chunk size cannot be greater than length(x) ($(N) > $(length(x)))"
+        @assert structural_length(x) >= N "chunk size cannot be greater than ForwardDiff.structural_length(x) ($(N) > $(structural_length(x)))"
 
         # precalculate loop bounds
-        xlen = length(x)
+        xlen = structural_length(x)
         remainder = xlen % N
         lastchunksize = ifelse(remainder == 0, N, remainder)
         lastchunkindex = xlen - lastchunksize + 1

src/prelude.jl

@@ -12,15 +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})
+structural_length(x::AbstractArray) = length(x)
+function structural_length(x::Union{LowerTriangular,UpperTriangular})
     n = size(x, 1)
     return (n * (n + 1)) >> 1
 end
-_length_structural_nonzero_indices(x::Diagonal) = size(x, 1)
+structural_length(x::Diagonal) = size(x, 1)
 
 function Chunk(x::AbstractArray, threshold::Integer = DEFAULT_CHUNK_THRESHOLD)
-    return Chunk(_length_structural_nonzero_indices(x), threshold)
+    return Chunk(structural_length(x), threshold)
 end
 
 # Constrained to `N <= threshold`, minimize (in order of priority):