Added Functionality for FieldVectors (from StaticArrays)

[kylejbrown17]

Hi Team,

This is my attempt to extend all SArray functionality to FieldVectors. This resolves issue #305. The basic approach mostly involves replacing SArray with Union{FieldVector, SArray} in method definitions. I have included tests to mirror those used for testing SArray codepaths for Gradients, Hessians and Jacobians.

[zsunberg]

@kylejbrown17 thanks! I added some more tests, and it looks like there is a stack overflow when calculating the jacobian. We should track this down and fix it before merging. Can you give me push access to your fork of this repo so that I can push to this PR? Thanks!

[zsunberg]

Oh, ok, I see the problem... Not sure if it is by design. The issue is in StaticArrays.jl. Will file an issue there. I fixed the tests, so this should be ready to merge.

[jrevels]

The type constraint on x allows this (and the other similar changes in this PR) to technically work, since we can probably assume that no upstream code will re-overload length or eltype for Union{FieldVector, SArray}.

However, it's a bit sketchy...it makes it easier for bugs to accidentally be introduced in the future via seemingly innocent changes (e.g. loosening the type constraint on x).

Would it be hard to just pull these out of the type parameter directly instead?

[kylejbrown17]

It seems like calling length and eltype makes this more straightforward to extend to additional types. @zsunberg and I decided to retrieve N and V via length and eltype because it was already quite complicated to try to extract them directly from the type parameters.

I've tried a few different variations on this:

@generated function dualize(::Type{T}, x::Union{FieldVector{N,V}, SArray{S,V,D,N}}) where {T,S,V,D,N}

Everything I've tried fails when called on an instance of FieldVector, although most of the approaches seem to work on SVector. I'm not 100% sure what you mean by "loosen the type constraint", but it seems like that would cause an even bigger headache if we want to extract the parameters this way.

If we tried to extract type parameters directly within the function body, we'd need a separate case for each type (because V and N are in different places for each type):

N,V = supertype(typeof(x)).parameters              # <: FieldVector
S,V,D,N = typeof(x).parameters             # SArray

Either of the above ways requires more customization than simply calling length and eltype. I think this way is more stable. It should work for vector types for which length and eltype are defined.

Another option is to have two different function definitions in every case where we need V and N.

[jrevels]

I think this way is more stable. It should work for vector types for which length and eltype are defined.

I think you might misunderstand my concern, which is my bad - I didn't explain the issue directly enough. If this were normal code, then of course I'd be fine with length and eltype getting used here.

However, this code is being called in an @generated thunk, and in general, calling overloadable methods inside of @generated function thunks is something that should be done with care, because downstream code can cause the generator to be "invalidated" in ways the compiler can't easily predict. Such code can cause world-age errors or silently incorrect method lookups, and are very hard for end-users to fix/debug/understand because they can depend on precompilation settings, code load order, etc. Thus, we want to be very careful about potentially introducing these kinds of bugs.

In this case, since FieldVector is an abstract type, so I could subtype it and overload length and eltype on that subtype. For the sake of example, let's call this subtype MyFieldVectorSubtype. If a user loaded ForwardDiff, called dualize on some other <:FieldVector, then loaded the code that defines FieldVectorSubtype, then called dualize on an instance of MyFieldVectorSubtype, the compiler could possibly throw a world-age error (most likely), silently call the wrong methods, run into some other compilation error, etc.

While this may seem like a niche class of bugs, they can crop up in practice quite regularly; for one ForwardDiff-related instance of such a bug, see JuliaArrays/StaticArrays.jl#316.

By the way, I really appreciate this PR, and am excited for this feature!

EDIT: With the above in mind, it might be sufficient to just define some internal helper functions, and leave a note in the code, e.g.

# these should not be overloaded for new types as doing so could 
# invalidate several of ForwardDiff's internal generated functions
_length(::Type{<:FieldVector{N}}) where {N} = N
_length(::Type{<:SArray{<:Any,<:Any,<:Any,N}}) where {N} = N
_eltype(::Type{<:FieldVector{<:Any,V}}) where {V} = V
_eltype(::Type{<:SArray{<:Any,V}}) where {V} = V

[zsunberg]

Ah ok, yeah, we didn't understand the specific concerns with @generated

What about

@generated function dualize(::Type{T}, x::Union{FieldVector{N,V}, SArray{<:Any,V,<:Any,N}}) where {T,V,N}

?

[zsunberg]

(we didn't know about the <:Any syntax in the type parameters)

[zsunberg]

If that doesn't work, I think two dualize functions would be a little simpler than _length and _eltype

EDIT: never mind - I didn't see that this is scattered throughout the other code

[kylejbrown17]

@jrevels, thanks for clarifying! This makes more sense. We'll incorporate the changes you suggested and get back to you.

[zsunberg]

We should delete this commented out line.

[kylejbrown17]

Ah, sorry about that. Not sure why I left that sitting in there.

[oschub]

I was interested in using FieldVectors in combination with ForwardDiff and asked in the Julia Slack channel for help since this PR has already been open for a long time. Mateusz Baran suggested the following:

ForwardDiff.dualize calls StaticArrays.similar_type at generation time, not call time, which means new methods of that function may or may not be picked (like the one you've added). This modification solves it:

@generated function ForwardDiff.dualize(::Type{T}, x::StaticArray) where T
    N = length(x)
    dx = Expr(:tuple, [:(ForwardDiff.Dual{T}(x[$i], chunk, Val{$i}())) for i in 1:N]...)
    return quote
        $(Expr(:meta, :inline))
        chunk = ForwardDiff.Chunk{$N}()
        V = StaticArrays.similar_type($x, $(ForwardDiff.Dual{T,eltype(x),N}))
        return V($(dx))
    end
end

I think this should be fine. I'd say my version is safer than the one currently present in ForwardDiff but both the current one and mine are a bit less safe then what was desired in that PR. I really doubt anyone would overload eltype or length for a subtype of FieldVector but that's actually where it's a bit unsafe.

To be honest these generated functions are a bit over my head but I wanted to share the suggestion from Mateusz Baran here to perhaps revive the discussion in this PR and to prevent the suggestion from getting lost due to Slack's message history limit.