work around length problem

Author: mcabbott

src/apiutils.jl

@@ -19,7 +19,7 @@ end
 ###################################
 
 @generated function dualize(::Type{T}, x::StaticArray) where T
-    N = length(x)
+    N = _static_length(StaticArraysCore.Size(x))
     dx = Expr(:tuple, [:(Dual{T}(x[$i], chunk, Val{$i}())) for i in 1:N]...)
     V = StaticArraysCore.similar_type(x, Dual{T,eltype(x),N})
     return quote
@@ -29,6 +29,9 @@ end
     end
 end
 
+# This works around length(::Type{StaticArray}) not being defined in this world-age:
+_static_length(::StaticArraysCore.Size{s}) where {s} = StaticArraysCore.tuple_prod(s)
+
 @inline static_dual_eval(::Type{T}, f, x::StaticArray) where T = f(dualize(T, x))
 
 function vector_mode_dual_eval!(f::F, cfg::Union{JacobianConfig,GradientConfig}, x) where {F}

src/gradient.jl

@@ -56,7 +56,8 @@ gradient(f, x::Real) = throw(DimensionMismatch("gradient(f, x) expects that x is
 #####################
 
 @generated function extract_gradient(::Type{T}, y::Real, x::S) where {T,S<:StaticArray}
-    result = Expr(:tuple, [:(partials(T, y, $i)) for i in 1:length(x)]...)
+    N = _static_length(StaticArraysCore.Size(S))
+    result = Expr(:tuple, [:(partials(T, y, $i)) for i in 1:N]...)
     return quote
         $(Expr(:meta, :inline))
         V = StaticArraysCore.similar_type(S, valtype($y))

src/jacobian.jl

@@ -96,8 +96,10 @@ jacobian(f, x::Real) = throw(DimensionMismatch("jacobian(f, x) expects that x is
 # result extraction #
 #####################
 
-@generated function extract_jacobian(::Type{T}, ydual::StaticArray, x::S) where {T,S<:StaticArray}
-    M, N = length(ydual), length(x)
+@generated function extract_jacobian(::Type{T}, ydual::Sy, x::S) where {T, Sy<:StaticArray, S<:StaticArray}
+    # M, N = length(ydual), length(x)
+    M = _static_length(StaticArraysCore.Size(Sy))
+    N = _static_length(StaticArraysCore.Size(S))
     result = Expr(:tuple, [:(partials(T, ydual[$i], $j)) for i in 1:M, j in 1:N]...)
     return quote
         $(Expr(:meta, :inline))